Merge branch 'pycolleff_integration' into develop

mbtrack2/impedance/wakefield.py

@@ -600,6 +600,8 @@ class Impedance(ComplexData):
         Return a WakeFunction object from the impedance data.
     plot()
         Plot the impedance data.
+    to_pycolleff()
+        Convenience method to export impedance to pycolleff.
     """
 
     def __init__(self,
@@ -780,6 +782,34 @@ class Impedance(ComplexData):
         label = r"$Z_{" + self.component_type + r"}$" + unit
         ax.set_ylabel(label)
         return ax
+    
+    def to_pycolleff(self):
+        """
+        Convenience method to export impedance to pycolleff.
+        Only implemented for longitudinal impedance.
+
+        Returns
+        -------
+        imp : pycolleff.longitudinal_equilibrium.ImpedanceSource
+            pycolleff ImpedanceSource object
+        """
+        # pycolleff is a soft dependency
+        from pycolleff.longitudinal_equilibrium import ImpedanceSource
+        # avoid circular import
+        from mbtrack2.utilities.misc import double_sided_impedance
+        imp = ImpedanceSource()
+        if self.component_type == "long":
+            double_sided_impedance(self)
+            # Negative imag sign convention !
+            imp.zl_table = self.data["real"].to_numpy() - 1j*self.data["imag"].to_numpy()
+            imp.ang_freq_table = self.data.index.to_numpy() * 2 * np.pi
+        else:
+            raise NotImplementedError()
+        
+        imp.calc_method = ImpedanceSource.Methods.ImpedanceDFT
+        imp.active_passive = ImpedanceSource.ActivePassive.Passive
+        
+        return imp
 
 
 class WakeField:

mbtrack2/tracking/rf.py

@@ -7,7 +7,7 @@ import matplotlib.patches as mpatches
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.legend_handler import HandlerPatch
-
+from functools import partial
 from mbtrack2.instability import lcbi_growth_rate
 from mbtrack2.tracking.element import Element
 
@@ -191,6 +191,8 @@ class CavityResonator():
         Force feedback update from current CavityResonator parameters.
     sample_voltage()
         Sample the voltage seen by a zero charge particle during an RF period.
+    to_pycolleff()
+        Convenience method to export a CavityResonator to pycolleff.
     
     References
     ----------
@@ -1141,6 +1143,49 @@ class CavityResonator():
                           ref_frame="beam")
 
         return pos, voltage_rec
+    
+    def to_pycolleff(self, Impedance=True):
+        """
+        Convenience method to export a CavityResonator to pycolleff.
+        
+        Parameters
+        ----------
+        Impedance : bool, optional
+            If True, export as impedance (i.e. ImpedanceSource.Methods.ImpedanceDFT).
+            If False, export as wake (i.e. ImpedanceSource.Methods.Wake).
+            Default is True.
+
+        Returns
+        -------
+        cav : pycolleff ImpedanceSource object
+
+        """
+        # pycolleff is a soft dependency
+        from pycolleff.longitudinal_equilibrium import ImpedanceSource
+        
+        cav = ImpedanceSource()
+        cav.harm_rf = self.m
+        cav.Q = self.QL
+        RoverQ = self.RL/self.QL
+        cav.shunt_impedance = RoverQ * cav.Q
+        cav.ang_freq_rf = self.ring.omega1
+        cav.ang_freq = cav.harm_rf * cav.ang_freq_rf
+        cav.detune_w = 2 * np.pi * self.detune
+        if self.Vg != 0:
+            cav.active_passive = ImpedanceSource.ActivePassive.Active
+            if Impedance:
+                raise NotImplementedError()
+            else:
+                cav.calc_method = ImpedanceSource.Methods.Wake         
+        else:
+            cav.active_passive = ImpedanceSource.ActivePassive.Passive
+            if Impedance:
+                cav.calc_method = ImpedanceSource.Methods.ImpedanceDFT
+            else:
+                cav.calc_method = ImpedanceSource.Methods.Wake     
+        return cav
+        
+
 
 
 class ProportionalLoop():

mbtrack2/tracking/synchrotron.py

@@ -116,6 +116,8 @@ class Synchrotron:
         single or multi-harmonic RF systems.
     to_pyat(Vrf)
         Return a pyAT simple_ring element from the Synchrotron element data.
+    to_pycolleff(I0, Vrf, bunch_number)
+        Return a pycolleff Ring element from the Synchrotron element data.
     """
 
     def __init__(self, h, optics, particle, **kwargs):
@@ -584,3 +586,49 @@ class Synchrotron:
                                      U0=self.U0,
                                      TimeLag=TimeLag)
         return at_simple_ring
+    
+    def to_pycolleff(self, I0, Vrf, bunch_number):
+        """
+        Return a pycolleff Ring element from the Synchrotron element data.
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+        Vrf : float
+            RF voltage in [V].
+        bunch_number : int
+            Total number of bunches filled.
+
+        Returns
+        -------
+        ring : pycolleff.colleff.Ring
+            pycolleff Ring object.
+
+        """
+        # pycolleff is a soft dependency
+        from pycolleff.colleff import Ring
+        
+        ring = Ring()
+        ring.version = 'from_mbtrack2'
+        ring.rf_freq = self.f1
+        ring.mom_comp = self.ac  # momentum compaction factor
+        ring.energy = self.E0  # energy [eV]
+        ring.tuney = self.tune[1]  # vertical tune
+        ring.tunex = self.tune[0]  # horizontal tune
+        ring.chromx = self.chro[0]  # horizontal chromaticity
+        ring.chromy = self.chro[1]  # vertical chromaticity
+        ring.harm_num = self.h  # harmonic Number
+        ring.num_bun = bunch_number  # number of bunches filled
+        ring.total_current = I0  # total current [A]
+        ring.sync_tune = self.synchrotron_tune(Vrf)  # synchrotron tune
+        ring.espread = self.sigma_delta
+        ring.bunlen = self.sigma_0*c  # [m]
+        ring.damptx = self.tau[0]  # [s]
+        ring.dampty = self.tau[1]  # [s]
+        ring.dampte = self.tau[2]  # [s]
+        ring.en_lost_rad = self.U0 # [eV]
+        ring.gap_voltage = Vrf  # [V]
+        
+        return ring
+

mbtrack2/utilities/beamloading.py

@@ -6,18 +6,14 @@ Module where the BeamLoadingEquilibrium class is defined.
 import matplotlib.pyplot as plt
 import numpy as np
 from scipy.constants import c
-from scipy.integrate import quad
+from scipy.integrate import trapz
 from scipy.optimize import root
-
+from mbtrack2.utilities.spectrum import gaussian_bunch
 
 class BeamLoadingEquilibrium():
     """Class used to compute beam equilibrium profile for a given storage ring 
     and a list of RF cavities of any harmonic. The class assumes an uniform 
-    filling of the storage ring. Based on an extension of [1].
-
-    [1] Venturini, M. (2018). Passive higher-harmonic rf cavities with general
-    settings and multibunch instabilities in electron storage rings.
-    Physical Review Accelerators and Beams, 21(11), 114404.
+    filling of the storage ring. Based on [1] which is an extension of [2].
 
     Parameters
     ----------
@@ -30,8 +26,20 @@ class BeamLoadingEquilibrium():
     PHI : list of form factor phase
     B1 : lower intergration boundary
     B2 : upper intergration boundary
+    N : int or float, optinal
+        Number of points used for longitudinal grid.
+        Default is 1000.
+    
+    References
+    ----------
+    [1] Gamelin, A., & Yamamoto, N. (2021). Equilibrium bunch density 
+    distribution with multiple active and passive RF cavities. 
+    IPAC'21 (MOPAB069).
+    [2] Venturini, M. (2018). Passive higher-harmonic rf cavities with general
+    settings and multibunch instabilities in electron storage rings.
+    Physical Review Accelerators and Beams, 21(11), 114404.
+    
     """
-
     def __init__(self,
                  ring,
                  cavity_list,
@@ -40,7 +48,8 @@ class BeamLoadingEquilibrium():
                  F=None,
                  PHI=None,
                  B1=-0.2,
-                 B2=0.2):
+                 B2=0.2,
+                 N=1000):
         self.ring = ring
         self.cavity_list = cavity_list
         self.I0 = I0
@@ -57,7 +66,10 @@ class BeamLoadingEquilibrium():
         self.B1 = B1
         self.B2 = B2
         self.mpi = False
-        self.__version__ = "1.0"
+        self.N = N
+        self.z0 = np.linspace(self.B1, self.B2, self.N)
+        self.tau0 = self.z0/c
+        self.rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0)/c
 
         # Define constants for scaled potential u(z)
         self.u0 = self.ring.U0 / (self.ring.ac * self.ring.sigma_delta**2 *
@@ -65,6 +77,11 @@ class BeamLoadingEquilibrium():
         self.ug = np.zeros((self.n_cavity, ))
         self.ub = np.zeros((self.n_cavity, ))
         self.update_potentials()
+                  
+    def update_rho(self):
+        uexp_vals = self.uexp(self.z0)
+        A = trapz(uexp_vals, x=self.z0)
+        self.rho0 = uexp_vals / A
 
     def update_potentials(self):
         """Update potentials with cavity and ring data."""
@@ -90,9 +107,7 @@ class BeamLoadingEquilibrium():
 
     def center_of_mass(self):
         """Return center of mass position in [s]"""
-        z0 = np.linspace(self.B1, self.B2, 1000)
-        rho = self.rho(z0)
-        CM = np.average(z0, weights=rho)
+        CM = np.average(self.z0, weights=self.rho0)
         return CM / c
 
     def u(self, z):
@@ -122,12 +137,12 @@ class BeamLoadingEquilibrium():
 
     def uexp(self, z):
         return np.exp(-1 * self.u(z))
-
+  
     def integrate_func(self, f, g):
         """Return Integral[f*g]/Integral[f] between B1 and B2"""
-        A = quad(lambda x: f(x) * g(x), self.B1, self.B2)
-        B = quad(f, self.B1, self.B2)
-        return A[0] / B[0]
+        A = trapz(f*g, x=self.z0)
+        B = trapz(f, x=self.z0)
+        return A / B
 
     def to_solve(self, x, CM=True):
         """System of non-linear equation to solve to find the form factors F
@@ -150,6 +165,7 @@ class BeamLoadingEquilibrium():
         else:
             self.F = x[::2]
         self.PHI = x[1::2]
+        self.update_rho()
 
         # Compute system
         if self.auto_set_MC_theta:
@@ -158,12 +174,12 @@ class BeamLoadingEquilibrium():
                 cavity = self.cavity_list[i]
                 res[2 *
                     i] = self.F[i] * np.cos(self.PHI[i]) - self.integrate_func(
-                        lambda y: self.uexp(y),
-                        lambda y: np.cos(self.ring.k1 * cavity.m * y))
+                        self.uexp(self.z0),
+                        np.cos(self.ring.k1 * cavity.m * self.z0))
                 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))
+                        self.uexp(self.z0),
+                        np.sin(self.ring.k1 * cavity.m * self.z0))
             # Factor 1e-8 or 1e12 for better convergence
             if CM is True:
                 res[self.n_cavity * 2] = self.center_of_mass() * 1e12
@@ -175,27 +191,17 @@ class BeamLoadingEquilibrium():
                 cavity = self.cavity_list[i]
                 res[2 *
                     i] = self.F[i] * np.cos(self.PHI[i]) - self.integrate_func(
-                        lambda y: self.uexp(y),
-                        lambda y: np.cos(self.ring.k1 * cavity.m * y))
+                        self.uexp(self.z0),
+                        np.cos(self.ring.k1 * cavity.m * self.z0))
                 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))
+                        self.uexp(self.z0),
+                        np.sin(self.ring.k1 * cavity.m * self.z0))
         return res
 
-    def rho(self, z):
-        """Return bunch equilibrium profile at postion z"""
-        A = quad(lambda y: self.uexp(y), self.B1, self.B2)
-        return self.uexp(z) / A[0]
-
-    def plot_rho(self, z1=None, z2=None):
-        """Plot the bunch equilibrium profile between z1 and z2"""
-        if z1 is None:
-            z1 = self.B1
-        if z2 is None:
-            z2 = self.B2
-        z0 = np.linspace(z1, z2, 1000)
-        plt.plot(z0 / c * 1e12, self.rho(z0))
+    def plot_rho(self):
+        """Plot the bunch equilibrium profile"""
+        plt.plot(self.tau0 * 1e12, self.rho0)
         plt.xlabel(r"$\tau$ [ps]")
         plt.title("Equilibrium bunch profile")
 
@@ -231,38 +237,22 @@ class BeamLoadingEquilibrium():
             Vtot += cavity.deltaVRF(z, self.I0, self.F[i], self.PHI[i])
         return Vtot
 
-    def plot_dV(self, z1=None, z2=None):
-        """Plot the derivative of RF system total voltage between z1 and z2"""
-        if z1 is None:
-            z1 = self.B1
-        if z2 is None:
-            z2 = self.B2
-        z0 = np.linspace(z1, z2, 1000)
-        plt.plot(z0, self.dV(z0))
+    def plot_dV(self):
+        """Plot the derivative of RF system total voltage"""
+        plt.plot(self.z0, self.dV(self.z0))
         plt.xlabel("z [m]")
         plt.ylabel("Total RF voltage (V)")
 
-    def plot_voltage(self, z1=None, z2=None):
+    def plot_voltage(self):
         """Plot the RF system total voltage between z1 and z2"""
-        if z1 is None:
-            z1 = self.B1
-        if z2 is None:
-            z2 = self.B2
-        z0 = np.linspace(z1, z2, 1000)
-        plt.plot(z0, self.voltage(z0))
+        plt.plot(self.z0, self.voltage(self.z0))
         plt.xlabel("z [m]")
         plt.ylabel("Total RF voltage (V)")
 
-    def std_rho(self, z1=None, z2=None):
+    def std_rho(self):
         """Return the rms bunch equilibrium size in [m]"""
-        if z1 is None:
-            z1 = self.B1
-        if z2 is None:
-            z2 = self.B2
-        z0 = np.linspace(z1, z2, 1000)
-        values = self.rho(z0)
-        average = np.average(z0, weights=values)
-        variance = np.average((z0 - average)**2, weights=values)
+        average = np.average(self.z0, weights=self.rho0)
+        variance = np.average((self.z0 - average)**2, weights=self.rho0)
         return np.sqrt(variance)
 
     def beam_equilibrium(self,
@@ -333,7 +323,7 @@ class BeamLoadingEquilibrium():
                 print("The algorithm has converged: " + str(sol.success))
 
         if plot:
-            self.plot_rho(self.B1 / 4, self.B2 / 4)
+            self.plot_rho()
 
         return sol
 
@@ -404,3 +394,19 @@ class BeamLoadingEquilibrium():
                 (2 * np.pi * HC.m**2 * self.F[1] * self.ring.h * I0 * f))
 
         return (eta, RQth, f)
+    
+    @property
+    def R_factor(self):
+        """
+        Touschek lifetime ratio as defined in [1].
+            
+        Reference
+        ---------
+        [1] : Byrd, J. M., and M. Georgsson. "Lifetime increase using passive 
+        harmonic cavities in synchrotron light sources." Physical Review 
+        Special Topics-Accelerators and Beams 4.3 (2001): 030701.
+
+        """
+        rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0)/c
+        R = trapz(rho0**2, self.tau0)/trapz(self.rho0**2, self.tau0)
+        return R