Improvement to WakeFunction and Impedance

Improve WakeFunction and Impedance methods to better take into account transverse wakes.
Update read_IW2D function to account for pandas library changes in the last versions.

mbtrack2/impedance/wakefield.py

@@ -273,7 +273,9 @@ class WakeFunction(ComplexData):
         self._wake_type = value
         
     def from_wakepotential(self, file_name, bunch_length, bunch_charge, 
-                           freq_lim, nout=None, trim=False):
+                           freq_lim, wake_type="long", divide_by=None, 
+                           nout=None, trim=False, axis0='dist', 
+                           axis0_scale=1e-3, axis1_scale=1e-12):
         """
         Compute a wake function from a wake potential file and load it to the 
         WakeFunction object.
@@ -288,6 +290,11 @@ class WakeFunction(ComplexData):
             Bunch charge in [C].
         freq_lim : float
             The maximum frequency for calculating the impedance [Hz].
+        wake_type : str, optional
+            Type of the wake: "long", "xdip", "xquad", ...
+        divide_by : float, optional
+            Divide the wake potential by a value. Mainly used to normalize 
+            transverse wake function by displacement.
         nout : int, optional
             Length of the transformed axis of the output. If None, it is taken
             to be 2*(n-1) where n is the length of the input. If nout > n, the
@@ -298,15 +305,27 @@ class WakeFunction(ComplexData):
             to zero where time<0. 
             If False, the original result is preserved.
             If a float is given, the pseudo wake function is trimmed from 
-            time <= trim to 0. 
+            time <= trim to 0.
+        axis0 : {'dist', 'time'}, optional
+            Viariable of the data file's first column. Use 'dist' for spacial 
+            distance. Otherwise, 'time' for temporal distance. 
+        axis0_scale : float, optional
+            Scale of the first column with respect to meter if axis0 = 'dist',
+            or to second if axis0 = 'time'.
+        axis1_scale : float, optional
+            Scale of the wake potential (in second column) with respect to V/C.
     
         """
         
         imp = Impedance()
         imp.from_wakepotential(file_name=file_name, bunch_length=bunch_length,
-                               bunch_charge=bunch_charge, freq_lim=freq_lim)
+                               bunch_charge=bunch_charge, freq_lim=freq_lim,
+                               impedance_type=wake_type, divide_by=divide_by,
+                               axis0=axis0, axis0_scale=axis0_scale,
+                               axis1_scale=axis1_scale)
         wf = imp.to_wakefunction(nout=nout, trim=trim)
-        self.__init__(variable=wf.data.index, function=wf.data["real"])
+        self.__init__(variable=wf.data.index, function=wf.data["real"],
+                      wake_type=wake_type)
         
 class Impedance(ComplexData):
     """
@@ -505,8 +524,8 @@ class Impedance(ComplexData):
         return kloss
     
     def from_wakepotential(self, file_name, bunch_length, bunch_charge, 
-                           freq_lim, axis0='dist', axis0_scale=1e-3, 
-                           axis1_scale=1e-12):
+                           freq_lim, impedance_type="long", divide_by=None,
+                           axis0='dist', axis0_scale=1e-3, axis1_scale=1e-12):
         """
         Compute impedance from wake potential data and load it to the Impedance
         object.
@@ -515,12 +534,17 @@ class Impedance(ComplexData):
         ----------
         file_name : str
             Text file that contains wake potential data.
-        freq_lim : float
-            The maximum frequency for calculating the impedance [Hz].
         bunch_length : float
             Electron bunch lenth [m]. 
         bunch_charge : float
             Total bunch charge [C].
+        freq_lim : float
+            The maximum frequency for calculating the impedance [Hz].
+        impedance_type : str, optional
+            Type of the impedance: "long", "xdip", "xquad", ...
+        divide_by : float, optional
+            Divide the impedance by a value. Mainly used to normalize transverse 
+            impedance by displacement.
         axis0 : {'dist', 'time'}, optional
             Viariable of the data file's first column. Use 'dist' for spacial 
             distance. Otherwise, 'time' for temporal distance. 
@@ -541,6 +565,8 @@ class Impedance(ComplexData):
             raise TypeError('Type of axis 0 not recognized.')
             
         wp = wp0 / axis1_scale
+        if divide_by is not None:
+            wp = wp / divide_by
         
         # FT of wake potential
         sampling = tau[1] - tau[0]
@@ -558,10 +584,15 @@ class Impedance(ComplexData):
         
         dft_rho_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
                                   1j*mu*2*np.pi*freq_trun)*bunch_charge
-        imp = dft_wake_trun/dft_rho_trun*-1*bunch_charge
-        
-        super().__init__(variable=freq_trun, function=imp)
-        self.data.index.name = "frequency [Hz]"    
+        if impedance_type == "long":
+            imp = dft_wake_trun/dft_rho_trun*-1*bunch_charge
+        elif (impedance_type == "xdip") or (impedance_type == "ydip"):
+            imp = dft_wake_trun/dft_rho_trun*-1j*bunch_charge
+        else:
+            raise NotImplementedError(impedance_type + " is not correct.")
+            
+        self.__init__(variable=freq_trun, function=imp, 
+                         impedance_type=impedance_type)
         
     def to_wakefunction(self, nout=None, trim=False):
         """
@@ -608,7 +639,8 @@ class Impedance(ComplexData):
             i_neg = np.where(time<trim)[0]
             Wlong[i_neg] = 0
                     
-        wf = WakeFunction(variable=time, function=Wlong)
+        wf = WakeFunction(variable=time, function=Wlong, 
+                          wake_type=self.impedance_type)
         return wf
     
 class WakeField:

mbtrack2/utilities/read_impedance.py

@@ -60,7 +60,7 @@ def read_IW2D(file, file_type='Zlong'):
         Data from file.
     """
     if file_type[0] == "Z":
-        df = pd.read_csv(file, sep = ' ', header = None, 
+        df = pd.read_csv(file, delim_whitespace=True, header = None, 
                          names = ["Frequency","Real","Imaginary"], skiprows=1)
         df.set_index("Frequency", inplace = True)
         df = df[df["Real"].notna()]
@@ -69,7 +69,7 @@ def read_IW2D(file, file_type='Zlong'):
                            function = df["Real"] + 1j*df["Imaginary"],
                            impedance_type=file_type[1:])
     elif file_type[0] == "W":
-        df = pd.read_csv(file, sep = ' ', header = None, 
+        df = pd.read_csv(file, delim_whitespace=True, header = None, 
                          names = ["Distance","Wake"], skiprows=1)
         df["Time"] = df["Distance"] / c
         df.set_index("Time", inplace = True)