Merge branch 'rf' into Prepare_version

tracking/monitors/monitors.py

@@ -1370,13 +1370,34 @@ class CavityMonitor(Monitor):
         group_name = cavity_name
         dict_buffer = {"cavity_phasor_record":(ring.h, buffer_size,),
                        "beam_phasor_record":(ring.h, buffer_size,),
-                       "detune":(buffer_size,)}
+                       "detune":(buffer_size,),
+                       "psi":(buffer_size,),
+                       "Vg":(buffer_size,),
+                       "theta_g":(buffer_size,),
+                       "Pg":(buffer_size,),
+                       "Rs":(buffer_size,),
+                       "Q":(buffer_size,),
+                       "QL":(buffer_size,)}
         dict_file = {"cavity_phasor_record":(ring.h, total_size,),
                      "beam_phasor_record":(ring.h, total_size,),
-                      "detune":(total_size,)}
+                     "detune":(total_size,),
+                     "psi":(total_size,),
+                     "Vg":(total_size,),
+                     "theta_g":(total_size,),
+                     "Pg":(total_size,),
+                     "Rs":(total_size,),
+                     "Q":(total_size,),
+                     "QL":(total_size,)}
         dict_dtype = {"cavity_phasor_record":complex,
                       "beam_phasor_record":complex,
-                      "detune":float}
+                      "detune":float,
+                      "psi":float,
+                      "Vg":float,
+                      "theta_g":float,
+                      "Pg":float,
+                      "Rs":float,
+                      "Q":float,
+                      "QL":float}
         
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode, 

tracking/monitors/plotting.py

@@ -1072,9 +1072,12 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
             - "turn" plots the phasor voltage and ange versus bunch index for
             a given turn.
             - "streak_volt" plots the phasor voltage versus bunch index and 
-            time
+            time.
             - "streak_angle" plots the phasor angle versus bunch index and 
-            time
+            time.
+            - "detune" or "psi" plots the detuning or tuning angle versus time.
+            - "power" plots the generator, cavity, beam and reflected power
+            versus time.
     bunch_number : int, optional
         Bunch number to select. The default is 0.
     turn : int, optional
@@ -1166,7 +1169,45 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         ax.set_xlabel("Bunch index")
         ax.set_ylabel("Number of turns")
         cbar = fig.colorbar(c, ax=ax)
-        cbar.set_label(ylabel) 
-    
+        cbar.set_label(ylabel)
+        
+    if plot_type == "detune" or plot_type == "psi":
+        
+        fig, ax = plt.subplots()
+        if plot_type == "detune":
+            data = np.array(cavity_data["detune"])*1e-3
+            ylabel = r"Detuning $\Delta f$ [kHz]"
+        elif plot_type == "psi":
+            data = np.array(cavity_data["psi"])
+            ylabel = r"Tuning angle $\psi$"
+            
+        ax.plot(time, data)
+        ax.set_xlabel("Number of turns")
+        ax.set_ylabel(ylabel)
+        
+    if plot_type == "power":
+        Vc = np.mean(np.abs(cavity_data["cavity_phasor_record"]),0)
+        theta = np.mean(np.angle(cavity_data["cavity_phasor_record"]),0)
+        try:
+            bunch_index = (file["Beam"]["current"][:,0] != 0).nonzero()[0]
+            I0 = np.nansum(file["Beam"]["current"][bunch_index,:],0)
+        except:
+            print("Beam monitor is needed to compute power.")
+            
+        Rs = np.array(cavity_data["Rs"])
+        Pc = Vc**2 / (2 * Rs)
+        Pb = I0 * Vc * np.cos(theta)
+        Pg = np.array(cavity_data["Pg"])
+        Pr = Pg - Pb - Pc
+        
+        fig, ax = plt.subplots()
+        ax.plot(time, Pg*1e-3, label="Generator power $P_g$ [kW]")
+        ax.plot(time, Pb*1e-3, label="Beam power $P_b$ [kW]")
+        ax.plot(time, Pc*1e-3, label="Dissipated cavity power $P_c$ [kW]")
+        ax.plot(time, Pr*1e-3, label="Reflected power $P_r$ [kW]")
+        ax.set_xlabel("Number of turns")
+        ax.set_ylabel("Power [kW]")
+        plt.legend()
+        
     file.close()
     return fig

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
@@ -173,6 +192,11 @@ class CavityResonator():
         self.beam_phasor = np.zeros(1, dtype=np.complex)
         self.beam_phasor_record = np.zeros((self.ring.h), dtype=np.complex)
         self.tracking = False
+        self.Vg = 0
+        self.theta_g = 0
+        self.Vgr = 0
+        self.theta_gr = 0
+        self.Pg = 0
         
     def init_tracking(self, beam):
         """
@@ -516,15 +540,34 @@ 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
 
     @property
     def Rs(self):
         """Shunt impedance [ohm]"""
-        return self._Rs
+        return self.Rs_per_cavity * self.Ncav
 
     @Rs.setter
     def Rs(self, value):
-        self._Rs = value
+        self.Rs_per_cavity = value / self.Ncav
 
     @property
     def Q(self):
@@ -600,7 +643,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 +732,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):
         """

vlasov/beamloading.py

@@ -26,8 +26,8 @@ class BeamLoadingEquilibrium():
     ring : Synchrotron object
     cavity_list : list of CavityResonator objects
     I0 : beam current in [A].
-    auto_set_MC_theta : if True, allow class to change generator phase for
-        CavityResonator objetcs with m = 1
+    auto_set_MC_theta : if True, allow class to change cavity phase for
+        CavityResonator objetcs with m = 1 (i.e. main cavities)
     F : list of form factor amplitude
     PHI : list of form factor phase
     B1 : lower intergration boundary
@@ -61,6 +61,10 @@ class BeamLoadingEquilibrium():
             * self.ring.E0 * self.ring.L)
         self.ug = np.zeros((self.n_cavity,))
         self.ub = np.zeros((self.n_cavity,))
+        self.update_potentials()
+            
+    def update_potentials(self):
+        """Update potentials with cavity and ring data."""
         for i in range(self.n_cavity):
             cavity = self.cavity_list[i]
             self.ug[i] = cavity.Vg / (
@@ -121,20 +125,24 @@ class BeamLoadingEquilibrium():
         B = quad(f, self.B1, self.B2)
         return A[0] / B[0]
 
-    def to_solve(self, x):
+    def to_solve(self, x, CM=True):
         """System of non-linear equation to solve to find the form factors F
         and PHI at equilibrum.
         The system is composed of Eq. (B6) and (B7) of [1] for each cavity.
-        If auto_set_MC_theta == True, the system also find the generator phase
-        to impose energy balance.
+        If auto_set_MC_theta == True, the system also find the main cavity 
+        phase to impose energy balance or cancel center of mass offset.
+        If CM is True, the system imposes zero center of mass offset,
+        if False, the system imposes energy balance.
         """
-        # Update values of F, PHI and theta_g
+        # Update values of F, PHI and theta
         if self.auto_set_MC_theta:
             self.F = x[:-1:2]
             for i in range(self.n_cavity):
                 cavity = self.cavity_list[i]
                 if cavity.m == 1:
-                    cavity.theta_g = x[-1]
+                    cavity.theta = x[-1]
+                    cavity.set_generator(0.5)
+                    self.update_potentials()
         else:
             self.F = x[::2]
         self.PHI = x[1::2]
@@ -148,8 +156,11 @@ class BeamLoadingEquilibrium():
                     lambda y: self.uexp(y), lambda y: np.cos(self.ring.k1 * cavity.m * y))
                 res[2 * i + 1] = self.F[i] * np.sin(self.PHI[i]) - self.integrate_func(
                     lambda y: self.uexp(y), lambda y: np.sin(self.ring.k1 * cavity.m * y))
-            # Factor 1e-8 for better convergence
-            res[self.n_cavity * 2] = self.energy_balance() * 1e-8
+            # Factor 1e-8 or 1e12 for better convergence
+            if CM is True:
+                res[self.n_cavity * 2] = self.center_of_mass() * 1e12
+            else:
+                res[self.n_cavity * 2] = self.energy_balance() * 1e-8
         else:
             res = np.zeros((self.n_cavity * 2,))
             for i in range(self.n_cavity):
@@ -242,7 +253,8 @@ class BeamLoadingEquilibrium():
         variance = np.average((z0 - average)**2, weights=values)
         return np.sqrt(variance)
 
-    def beam_equilibrium(self, x0=None, tol=1e-4, method='hybr', options=None, plot = False):
+    def beam_equilibrium(self, x0=None, tol=1e-4, method='hybr', options=None, 
+                         plot = False, CM=True):
         """Solve system of non-linear equation to find the form factors F
         and PHI at equilibrum. Can be used to compute the equilibrium bunch
         profile.
@@ -254,6 +266,8 @@ class BeamLoadingEquilibrium():
         method : method used by scipy.optimize.root to solve the system
         options : options given to scipy.optimize.root
         plot : if True, plot the equilibrium bunch profile
+        CM : if True, the system imposes zero center of mass offset,
+        if False, the system imposes energy balance.
         
         Returns
         -------
@@ -262,12 +276,17 @@ class BeamLoadingEquilibrium():
         if x0 is None:
             x0 = [1, 0] * self.n_cavity
             if self.auto_set_MC_theta:
-                x0 = x0 + [self.cavity_list[0].theta_g]
+                x0 = x0 + [self.cavity_list[0].theta]
 
-        print("The initial energy balance is " +
-              str(self.energy_balance()) + " eV")
+        if CM:
+            print("The initial center of mass offset is " +
+                  str(self.center_of_mass()*1e12) + " ps")
+        else:
+            print("The initial energy balance is " +
+                  str(self.energy_balance()) + " eV")
 
-        sol = root(self.to_solve, x0, tol=tol, method=method, options=options)
+        sol = root(lambda x : self.to_solve(x, CM), x0, tol=tol, method=method, 
+                   options=options)
 
         # Update values of F, PHI and theta_g
         if self.auto_set_MC_theta:
@@ -275,13 +294,17 @@ class BeamLoadingEquilibrium():
             for i in range(self.n_cavity):
                 cavity = self.cavity_list[i]
                 if cavity.m == 1:
-                    cavity.theta_g = sol.x[-1]
+                    cavity.theta = sol.x[-1]
         else:
             self.F = sol.x[::2]
         self.PHI = sol.x[1::2]
 
-        print("The final energy balance is " +
-              str(self.energy_balance()) + " eV")
+        if CM:
+            print("The final center of mass offset is " +
+                  str(self.center_of_mass()*1e12) + " ps")
+        else:
+            print("The final energy balance is " +
+                  str(self.energy_balance()) + " eV")
         print("The algorithm has converged: " + str(sol.success))
         
         if plot: