Improve CavityResonator

Add Ncav parameter and Rs_per_cavity attribute to remove ambiguity about what should be used
Add attributes and methods to compute Pc, Pb and Pr.
Add method to set optimal coupling value.

tracking/rf.py

@@ -55,6 +55,7 @@ class RFCavity(Element):
 class CavityResonator():
     """Cavity resonator class for active or passive RF cavity with beam
     loading or HOM, based on [1].
+    Use cosine definition.
     
     Parameters
     ----------
@@ -62,13 +63,17 @@ class CavityResonator():
     m : int or float
         Harmonic number of the cavity.
     Rs : float
-        Shunt impedance of the cavity in [Ohm], defined as 0.5*Vc*Vc/Pc.
+        Shunt impedance of the cavities in [Ohm], defined as 0.5*Vc*Vc/Pc.
+        If Ncav = 1, used for the total shunt impedance.
+        If Ncav > 1, used for the shunt impedance per cavity.
     Q : float
         Quality factor of the cavity.
     QL : float
         Loaded quality factor of the cavity.
     detune : float
         Detuing of the cavity in [Hz], defined as (fr - m*ring.f1).
+    Ncav : int, optional
+        Number of cavities.
     Vc : float, optinal
         Total cavity voltage in [V].
     theta : float, optional
@@ -92,6 +97,8 @@ class CavityResonator():
         Cavity total phase in [rad].
     loss_factor : float
         Cavity loss factor in [V/C].
+    Rs_per_cavity : float
+        Shunt impedance of a single cavity in [Ohm], defined as 0.5*Vc*Vc/Pc.
     beta : float
         Coupling coefficient of the cavity.
     fr : float
@@ -102,6 +109,8 @@ class CavityResonator():
         Tuning angle in [rad].
     filling_time : float
         Cavity filling time in [s].
+    Pc : float
+        Power dissipated in the cavity walls in [W].
     Pg : float
         Generator power in [W].
     Vgr : float
@@ -126,8 +135,14 @@ class CavityResonator():
         Return beam voltage at resonance in [V].
     Vb(I0)
         Return beam voltage in [V].
+    Pb(I0)
+        Return power transmitted to the beam in [W].
+    Pr(I0)
+        Return power reflected back to the generator in [W].
     Z(f)
         Cavity impedance in [Ohm] for a given frequency f in [Hz].
+    set_optimal_coupling(I0)
+        Set coupling to optimal value.
     set_optimal_detune(I0)
         Set detuning to optimal conditions.
     set_generator(I0)
@@ -161,10 +176,14 @@ class CavityResonator():
     factories. In Frontiers of Particle Beams: Factories with e+ e-Rings 
     (pp. 293-311). Springer, Berlin, Heidelberg.
     """
-    def __init__(self, ring, m, Rs, Q, QL, detune, Vc=0, theta=0):
+    def __init__(self, ring, m, Rs, Q, QL, detune, Ncav=1, Vc=0, theta=0):
         self.ring = ring
         self.m = m
-        self.Rs = Rs
+        self.Ncav = Ncav
+        if Ncav != 1:
+            self.Rs_per_cavity = Rs
+        else:
+            self.Rs = Rs
         self.Q = Q
         self.QL = QL
         self.detune = detune
@@ -516,6 +535,26 @@ class CavityResonator():
     @m.setter
     def m(self, value):
         self._m = value
+        
+    @property
+    def Ncav(self):
+        """Number of cavities"""
+        return self._Ncav
+
+    @Ncav.setter
+    def Ncav(self, value):
+        self._Ncav = value
+        
+    @property
+    def Rs_per_cavity(self):
+        """Shunt impedance of a single cavity in [Ohm], defined as 
+        0.5*Vc*Vc/Pc."""
+        return self._Rs_per_cavity
+
+    @Rs_per_cavity.setter
+    def Rs_per_cavity(self, value):
+        self._Rs_per_cavity = value
+        self._Rs = self.Rs_per_cavity * self.Ncav
 
     @property
     def Rs(self):
@@ -524,7 +563,7 @@ class CavityResonator():
 
     @Rs.setter
     def Rs(self, value):
-        self._Rs = value
+        self.Rs_per_cavity = value / self.Ncav
 
     @property
     def Q(self):
@@ -600,7 +639,46 @@ class CavityResonator():
     def filling_time(self):
         """Cavity filling time in [s]"""
         return 2*self.QL/self.wr
-        
+    
+    @property
+    def Pc(self):
+        """Power dissipated in the cavity walls in [W]"""
+        return self.Vc**2 / (2 * self.Rs)
+    
+    def Pb(self, I0):
+        """
+        Return power transmitted to the beam in [W] - near Eq. (4.2.3) in [1].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        float
+            Power transmitted to the beam in [W].
+
+        """
+        return I0 * self.Vc * np.cos(self.theta)
+    
+    def Pr(self, I0):
+        """
+        Power reflected back to the generator in [W].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        float
+            Power reflected back to the generator in [W].
+
+        """
+        return self.Pg - self.Pb(I0) - self.Pc
+
     def Vbr(self, I0):
         """
         Return beam voltage at resonance in [V].
@@ -650,6 +728,19 @@ class CavityResonator():
 
         """
         self.psi = np.arctan(-self.Vbr(I0)/self.Vc*np.sin(self.theta))
+        
+    def set_optimal_coupling(self, I0):
+        """
+        Set coupling to optimal value - Eq. (4.2.3) in [1].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        """
+        self.beta = 1 + (2 * I0 * self.Rs * np.cos(self.theta) / 
+                         self.Vc)
                 
     def set_generator(self, I0):
         """