Various improvements

WakePotential -> add method to plot wake potential for a perfect gaussian bunch
utilities -> new function gaussian_bunch which gives the gaussian bunch profile
utilities -> read_IW2D now also reads wake function from IW2D
utilities -> new function read_IW2D_folder which reads the full results folder of IW2D into a WakeField element.
resistive_wall -> change the sign of the wake functions to agree with the code convention

collective_effects/__init__.py

@@ -13,5 +13,6 @@ from mbtrack2.collective_effects.wakefield import ComplexData, Impedance, WakeFu
 from mbtrack2.collective_effects.utilities import read_CST, read_IW2D, spectral_density
 from mbtrack2.collective_effects.utilities import gaussian_bunch_spectrum, beam_spectrum, effective_impedance
 from mbtrack2.collective_effects.utilities import yokoya_elliptic, beam_loss_factor
-from mbtrack2.collective_effects.utilities import double_sided_impedance
+from mbtrack2.collective_effects.utilities import double_sided_impedance, read_IW2D_folder
+from mbtrack2.collective_effects.utilities import gaussian_bunch
 from mbtrack2.collective_effects.impedance_model import ImpedanceModel
\ No newline at end of file

collective_effects/resistive_wall.py

@@ -155,7 +155,7 @@ class CircularResistiveWall(WakeField):
             ind = np.isnan(wl)
             wl[ind] = 0
     
-        return wl
+        return -1*wl
     
     def TransverseWakeFunction(self, time, exact=False):
         """
@@ -206,7 +206,7 @@ class CircularResistiveWall(WakeField):
                   / time**(1/2) * self.length * -1)
             ind = np.isnan(wt)
             wt[ind] = 0
-        return wt
+        return -1*wt
         
     def function(self, t, x):
         return ( (x**2 * np.exp(-1* (x**2) * t / self.t0) ) / (x**6 + 8) )

collective_effects/utilities.py

@@ -9,8 +9,10 @@ collective_effects module.
 
 import pandas as pd
 import numpy as np
+from pathlib import Path
 from scipy.interpolate import interp1d
-from mbtrack2.collective_effects.wakefield import Impedance
+from scipy.constants import c
+from mbtrack2.collective_effects.wakefield import Impedance, WakeFunction, WakeField
 
 def read_CST(file, impedance_type='long', divide_by=None):
     """
@@ -45,30 +47,82 @@ def read_CST(file, impedance_type='long', divide_by=None):
                        impedance_type=impedance_type)
     return result
 
-def read_IW2D(file, impedance_type='long'):
+def read_IW2D(file, file_type='Zlong'):
     """
-    Read IW2D file format into an Impedance object.
+    Read IW2D file format into an Impedance object or a WakeField object.
     
     Parameters
     ----------
     file : str
         Path to the file to read.
-    impedance_type : str, optional
-        Type of the Impedance object.
+    file_type : str, optional
+        Type of the Impedance or WakeField object.
         
     Returns
     -------
-    result : Impedance object
+    result : Impedance or WakeField object
         Data from file.
     """
+    if file_type[0] == "Z":
         df = pd.read_csv(file, sep = ' ', header = None, 
                          names = ["Frequency","Real","Imaginary"], skiprows=1)
         df.set_index("Frequency", inplace = True)
+        df = df[df["Real"].notna()]
+        df = df[df["Imaginary"].notna()]
         result = Impedance(variable = df.index,
                            function = df["Real"] + 1j*df["Imaginary"],
-                       impedance_type=impedance_type)
+                           impedance_type=file_type[1:])
+    elif file_type[0] == "W":
+        df = pd.read_csv(file, sep = ' ', header = None, 
+                         names = ["Distance","Wake"], skiprows=1)
+        df["Time"] = df["Distance"] / c
+        df.set_index("Time", inplace = True)
+        if np.any(df.isna()):
+            index = df.isna().values
+            df = df.interpolate()
+            print("Nan values have been interpolated to:")
+            print(df[index])
+        # if file_type == "Wlong":
+        #     df["Wake"] = df["Wake"]*-1
+        result = WakeFunction(variable = df.index,
+                           function = df["Wake"],
+                           wake_type=file_type[1:])
+    else:
+        raise ValueError("file_type should begin by Z or W.")
     return result
 
+def read_IW2D_folder(folder, suffix):
+    """
+    Read IW2D results into a WakeField object.
+    
+    Parameters
+    ----------
+    file : str
+        Path to the file to read.
+        
+    Returns
+    -------
+    result : WakeField object
+        WakeField object with Impedance and WakeFunction objects from the 
+        different files.
+    """
+    types = {"W" : WakeFunction,
+             "Z" : Impedance}
+    components = ["long", "xdip", "ydip", "xquad", "yquad"]
+    
+    data_folder = Path(folder)
+    
+    list_for_wakefield = []
+    for key, item in types.items():
+        for component in components:
+            name = data_folder / (key + component + suffix)
+            res = read_IW2D(file=name, file_type=key + component)
+            list_for_wakefield.append(res)
+            
+    wake = WakeField(list_for_wakefield)
+    
+    return wake
+
 def spectral_density(frequency, sigma, m = 1, mode="Hermite"):
     """
     Compute the spectral density of different modes for various values of the
@@ -124,6 +178,24 @@ def gaussian_bunch_spectrum(frequency, sigma):
     """
     return np.exp(-1/2*(2*np.pi*frequency)**2*sigma**2)
 
+def gaussian_bunch(time, sigma): 
+    """
+    Compute a Gaussian bunch profile.
+
+    Parameters
+    ----------
+    time : array
+        sample points of the bunch profile in [s].
+    sigma : float
+        RMS bunch length in [s].
+
+    Returns
+    -------
+    bunch_profile : array
+        Bunch profile in [s**-1] sampled at points given in time.
+    """
+    return np.exp(-1/2*(time**2/sigma**2))/(sigma*np.sqrt(2*np.pi))
+        
         
 def beam_spectrum(frequency, M, bunch_spacing, sigma=None, 
                   bunch_spectrum=None): 

collective_effects/wakefield.py

@@ -361,6 +361,7 @@ class Impedance(ComplexData):
             component_coefficients = table[self.impedance_type].to_dict()
         except KeyError:
             print('Impedance type {} does not exist'.format(self.impedance_type))
+            raise
         
         self.a = component_coefficients["a"]
         self.b = component_coefficients["b"]

tracking/wakepotential.py

@@ -13,6 +13,7 @@ import pandas as pd
 from scipy import signal
 from scipy.interpolate import interp1d
 from mbtrack2.tracking.element import Element
+from mbtrack2.collective_effects.utilities import gaussian_bunch
    
 class WakePotential(Element):
     """
@@ -166,7 +167,7 @@ class WakePotential(Element):
         return dipole0
     
     
-    def prepare_wakefunction(self, wake_type):
+    def prepare_wakefunction(self, wake_type, tau, save_data=True):
         """
         Prepare the wake function of a given wake_type to be used for the wake
         potential computation. 
@@ -179,6 +180,10 @@ class WakePotential(Element):
         ----------
         wake_type : str
             Type of the wake function to prepare: "Wlong", "Wxdip", ...
+        tau : array
+            Time domain array of the bunch profile in [s].
+        save_data : bool, optional
+            If True, the results are saved as atributes.
 
         Returns
         -------
@@ -202,7 +207,7 @@ class WakePotential(Element):
             W0 = np.array(getattr(self.wakefield, wake_type).data["real"])
         
         # Keep only the wake function on the rho window
-        ind = np.all([min(self.tau[0], 0) < tau0, max(self.tau[-1], 0) > tau0],
+        ind = np.all([min(tau[0], 0) < tau0, max(tau[-1], 0) > tau0],
                      axis=0)
         tau0 = tau0[ind]
         W0 = W0[ind]
@@ -229,9 +234,9 @@ class WakePotential(Element):
                 W0 = np.insert(W0, 0, np.zeros(n_to_add))
                 
         # Check is the wf is shorter than rho then add zeros
-        if (tau0[0] > self.tau[0]) or (tau0[-1] < self.tau[-1]):
-            n = max(int(np.ceil((tau0[0] - self.tau[0])/dtau0)),
-                    int(np.ceil((self.tau[-1] - tau0[-1])/dtau0)))
+        if (tau0[0] > tau[0]) or (tau0[-1] < tau[-1]):
+            n = max(int(np.ceil((tau0[0] - tau[0])/dtau0)),
+                    int(np.ceil((tau[-1] - tau0[-1])/dtau0)))
             
             tau0 = np.arange(tau0[0] - dtau0*n, 
                              tau0[-1] + dtau0*(n+1), 
@@ -240,9 +245,11 @@ class WakePotential(Element):
             n_to_add = len(tau0) - len(W0)
             W0 = np.insert(W0, len(W0), np.zeros(n_to_add))
 
+        if save_data:
             setattr(self, "tau0_" + wake_type, tau0)
             setattr(self, "dtau0_" + wake_type, dtau0)
             setattr(self, "W0_" + wake_type, W0)
+            
         return (tau0, dtau0, W0)
         
     def get_wakepotential(self, bunch, wake_type):
@@ -263,7 +270,7 @@ class WakePotential(Element):
 
         """
 
-        (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type)
+        (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type, self.tau)
         
         interp_profile = interp1d(self.tau, self.rho, fill_value=0, 
                                      bounds_error=False)
@@ -375,6 +382,106 @@ class WakePotential(Element):
             
         return fig
     
+    def get_gaussian_wakepotential(self, sigma, wake_type, dipole=1e-3):
+        """
+        Return the wake potential computed using a perfect gaussian profile.
+
+        Parameters
+        ----------
+        sigma : float
+            RMS bunch length in [s].
+        wake_type : str
+            Wake function type: "Wlong", "Wxdip", ...
+        dipole : float, optional
+            Dipole moment to consider in [m], (uniform dipole moment).
+        Returns
+        -------
+        tau0 : array
+            Time base in [s].
+        W0 : array
+            Wake function.
+        Wp : array
+            Wake potential.
+        profile0 : array
+            Gaussian bunch profile.
+        dipole0 : array
+            Dipole moment.
+
+        """
+        
+        tau = np.linspace(-10*sigma,10*sigma, int(1e3))
+        (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type, tau, False)
+        
+        profile0 = gaussian_bunch(tau0, sigma)
+        dipole0 = np.ones_like(profile0)*dipole
+        
+        if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
+            Wp = signal.convolve(profile0, W0*-1, mode='same')*dtau0
+        elif wake_type == "Wxdip":
+            Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
+        elif wake_type == "Wydip":
+            Wp = signal.convolve(profile0*dipole0, W0*-1, mode='same')*dtau0
+        else:
+            raise ValueError("This type of wake is not taken into account.")
+
+        return tau0, W0, Wp, profile0, dipole0
+        
+    def plot_gaussian_wake(self, sigma, wake_type, dipole=1e-3, plot_rho=True, 
+                           plot_dipole=False, plot_wake_function=True):
+        """
+        Plot the wake potential of a given type for a perfect gaussian bunch.
+
+        Parameters
+        ----------
+        sigma : float
+            RMS bunch length in [s].
+        wake_type : str
+            Type of the wake to plot: "Wlong", "Wxdip", ...
+        dipole : float, optional
+            Dipole moment to consider in [m], (uniform dipole moment).
+        plot_rho : bool, optional
+            Plot the normalised bunch profile. The default is True.
+        plot_dipole : bool, optional
+            Plot the normalised dipole moment. The default is False.
+        plot_wake_function : bool, optional
+            Plot the normalised wake function. The default is True.
+
+        Returns
+        -------
+        fig : figure
+
+        """
+        
+        labels = {"Wlong" : r"$W_{p,long}$ (V/pC)", 
+                  "Wxdip" : r"$W_{p,x}^{D} (V/pC)$",
+                  "Wydip" : r"$W_{p,y}^{D} (V/pC)$",
+                  "Wxquad" : r"$W_{p,x}^{Q} (V/pC/m)$",
+                  "Wyquad" : r"$W_{p,y}^{Q} (V/pC/m)$"}
+        
+        tau0, W0, Wp, profile0, dipole0 = self.get_gaussian_wakepotential(sigma, wake_type, dipole)
+        
+        fig, ax = plt.subplots()
+        ax.plot(tau0*1e12, Wp*1e-12, label=labels[wake_type])
+        ax.set_xlabel("$\\tau$ (ps)")
+        ax.set_ylabel(labels[wake_type])
+
+        if plot_rho is True:
+            profile_rescaled = profile0/max(profile0)*max(np.abs(Wp))
+            ax.plot(tau0*1e12, profile_rescaled*1e-12, label=r"$\rho$ (a.u.)")
+            plt.legend()
+            
+        if plot_wake_function is True:
+            W0_rescaled = W0/max(W0)*max(np.abs(Wp))
+            ax.plot(tau0*1e12, W0_rescaled*1e-12, label=r"$W_{function}$ (a.u.)")
+            plt.legend()
+            
+        if plot_dipole is True:
+            dipole_rescaled = dipole0/max(dipole0)*max(np.abs(Wp))
+            ax.plot(tau0*1e12, dipole_rescaled*1e-12, label=r"Dipole moment (a.u.)")
+            plt.legend()
+            
+        return fig
+    
     def reference_loss(self, bunch):
         """
         Calculate the loss factor and kick factor from the wake potential and