Tools for impedance models

Changes in loss_factor to use frequency instead of omega
Read_CST forces to have positive real part of longitudinal impedances
New functions : Gaussian_bunch_spectrum, beam_spectrum, beam_loss_factor
ImpedanceModel.plot_area in GHz

collective_effects/tools.py

@@ -9,6 +9,7 @@ Tools and utilities functions for the Wakefield and Impedances classes.
 import pandas as pd
 import numpy as np
 from scipy.integrate import trapz
+from scipy.interpolate import interp1d
 from collective_effects.wakefield import Wakefield, Impedance, WakeFunction
 
 def read_CST(file, impedance_type='long', divide_by=None):
@@ -36,6 +37,8 @@ def read_CST(file, impedance_type='long', divide_by=None):
     if divide_by is not None:
         df["Real"] = df["Real"]/divide_by
         df["Imaginary"] = df["Imaginary"]/divide_by
+    if impedance_type == "long":
+        df["Real"] = np.abs(df["Real"])
     df.set_index("Frequency", inplace = True)
     result = Impedance(variable = df.index,
                        function = df["Real"] + 1j*df["Imaginary"],
@@ -74,32 +77,36 @@ def loss_factor(wake, sigma):
     
     Parameters
     ----------
-    wake: Wakefield, Impedance or WakeFunction object.
-    sigma: RMS bunch length in [s]
+    wake : Wakefield, Impedance or WakeFunction object.
+    sigma : float
+        RMS bunch length in [s]
     
     Returns
     -------
-    kloss: loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
-    impedance type.
+    kloss: float
+        Loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
+        impedance type.
         
-    [1] Handbook of accelerator physics and engineering, 3rd printing.
+    References
+    ----------
+    [1] : Handbook of accelerator physics and engineering, 3rd printing.
     """
     
     if isinstance(wake, Impedance):
         
         positive_index = wake.data.index > 0
-        omega = 2*np.pi*wake.data.index[positive_index]
-        gaussian_spectrum = np.exp(-1*omega**2*sigma**2)
+        frequency = wake.data.index[positive_index]
+        sd = spectral_density(frequency, sigma, m=0)
         
         if(wake.impedance_type == "long"):
-            kloss = trapz(x = omega, 
-                          y = 1/np.pi*wake.data["real"][positive_index]*gaussian_spectrum)
+            kloss = trapz(x = frequency, 
+                          y = 2*wake.data["real"][positive_index]*sd)
         elif(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"):
-            kloss = trapz(x = omega, 
-                          y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
+            kloss = trapz(x = frequency, 
+                          y = 2*wake.data["imag"][positive_index]*sd)
         elif(wake.impedance_type == "xquad" or wake.impedance_type == "yquad"):
-            kloss = trapz(x = omega, 
-                          y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
+            kloss = trapz(x = frequency, 
+                          y = 2*wake.data["imag"][positive_index]*sd)
         else:
             raise TypeError("Impedance type not recognized.")
 
@@ -133,8 +140,10 @@ def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
     Returns
     -------
     numpy array
-        
-    [1] Handbook of accelerator physics and engineering, 3rd printing.
+    
+    References
+    ----------
+    [1] : Handbook of accelerator physics and engineering, 3rd printing.
     """
     
     if mode == "Hermite":
@@ -142,6 +151,70 @@ def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
                 -1*(2*np.pi*frequency*sigma)**2)
     else:
         raise NotImplementedError("Not implemanted yet.")
+        
+def Gaussian_bunch_spectrum(frequency, sigma): 
+    """
+    Compute a Gaussian bunch spectrum [1].
+
+    Parameters
+    ----------
+    frequency : array
+        sample points of the beam spectrum in [Hz].
+    sigma : float
+        RMS bunch length in [s].
+
+    Returns
+    -------
+    bunch_spectrum : array
+        Bunch spectrum sampled at points given in frequency.
+        
+    References
+    ----------
+    [1] : Gamelin, A. (2018). Collective effects in a transient microbunching 
+    regime and ion cloud mitigation in ThomX. p86, Eq. 4.19
+    """
+    return np.exp(-1/2*(2*np.pi*frequency)**2*sigma**2)
+        
+        
+def beam_spectrum(frequency, M, bunch_spacing, sigma=None, 
+                  bunch_spectrum=None): 
+    """
+    Compute the beam spectrum assuming constant spacing between bunches [1].
+
+    Parameters
+    ----------
+    frequency : list or numpy array
+        sample points of the beam spectrum in [Hz].
+    M : int
+        Number of bunches.
+    bunch_spacing : float
+        Time between two bunches in [s].
+    sigma : float, optional
+        If bunch_spectrum is None then a Gaussian bunch with sigma RMS bunch 
+        length in [s] is assumed.
+    bunch_spectrum : array, optional
+        Bunch spectrum sampled at points given in frequency.
+
+    Returns
+    -------
+    beam_spectrum : array
+
+    References
+    ----------
+    [1] Rumolo, G - Beam Instabilities - CAS - CERN Accelerator School: 
+        Advanced Accelerator Physics Course - 2014, Eq. 9
+    """
+    
+    if bunch_spectrum is None:
+        bunch_spectrum = Gaussian_bunch_spectrum(frequency, sigma)
+        
+    beam_spectrum = (bunch_spectrum * np.exp(1j*np.pi*frequency * 
+                                             bunch_spacing*(M-1)) * 
+                     np.sin(M*np.pi*frequency*bunch_spacing) / 
+                     np.sin(np.pi*frequency*bunch_spacing))
+    
+    return beam_spectrum
+    
     
 def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
     """
@@ -279,3 +352,58 @@ def Yokoya_elliptic(x_radius , y_radius):
         yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])        
 
     return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)
+
+def beam_loss_factor(impedance, frequency, spectrum, ring):
+    """
+    Compute "beam" loss factor using the beam spectrum, uses a sum instead of 
+    integral compared to loss_factor [1].
+
+    Parameters
+    ----------
+    impedance : Impedance of type "long"
+    frequency : array
+        Sample points of spectrum.
+    spectrum : array
+        Beam spectrum to consider.
+    ring : Synchrotron object
+
+    Returns
+    -------
+    kloss_beam : float
+        Beam loss factor in [V/C].
+        
+    References
+    ----------
+    [1] : Handbook of accelerator physics and engineering, 3rd printing. 
+        Eq (3) p239.
+    """
+    pmax = np.floor(impedance.data.index.max()/ring.f0)
+    pmin = np.floor(impedance.data.index.min()/ring.f0)
+    
+    if pmin >= 0:
+        # Add negative part
+        negative_index = impedance.data.index*-1
+        negative_data = impedance.data.set_index(negative_index)
+        negative_data["imag"] = -1*negative_data["imag"]
+        negative_data = negative_data.drop(0)
+        
+        all_data = impedance.data.append(negative_data)
+        all_data = all_data.sort_index()
+        impedance.data = all_data
+        
+        pmin = -1*pmax
+    
+    p = np.arange(pmin+1,pmax)    
+    pf0 = p*ring.f0
+    ReZ = np.real(impedance(pf0))
+    spectral_density = np.abs(spectrum)**2
+    # interpolation of the spectrum is needed to avoid problems liked to 
+    # division by 0
+    # computing the spectrum directly to the frequency points gives
+    # wrong results
+    spect = interp1d(frequency, spectral_density)
+    kloss_beam = ring.f0 * np.sum(ReZ*spect(pf0))
+    
+    return kloss_beam
+    
+    

collective_effects/wakefield.py

@@ -694,7 +694,7 @@ class ImpedanceModel(Element):
             attr = attr_list[index]
             # Set all impedances with common indexes using + zero_impedance
             sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
-            ax.fill_between(sum_imp.data.index, total_imp, 
+            ax.fill_between(sum_imp.data.index*1e-9, total_imp, 
                             total_imp + sum_imp.data[component])
             total_imp += sum_imp.data[component]
             if attr_list is None:
@@ -705,16 +705,18 @@ class ImpedanceModel(Element):
         if sigma is not None:
             spect = tools.spectral_density(zero_impedance.data.index, sigma)
             spect = spect/spect.max()*total_imp.max()
-            ax.plot(sum_imp.data.index, spect, 'r', linewidth=2.5)
+            ax.plot(sum_imp.data.index*1e-9, spect, 'r', linewidth=2.5)
         
-        ax.legend(legend)            
-        ax.set_xlabel("Frequency [Hz]")
+        ax.legend(legend, loc="upper left")            
+        ax.set_xlabel("Frequency [GHz]")
         
         if Z_type == "Zlong":
             ax.set_ylabel(Z_type + " [Ohm] - " + component + " part")
         else:
             ax.set_ylabel(Z_type + " [Ohm/m] - " + component + " part")
             
+        return fig
+            
     def summary(self,  sigma, attr_list=None, list_components=None):
         
         if attr_list is None: