Add documentation and rename parameters

Add docstrings to Impedance and WakeFunction class
Rename functions for deconvolution method

collective_effects/utilities.py

@@ -9,7 +9,6 @@ collective_effects module.
 
 import pandas as pd
 import numpy as np
-import matplotlib.pyplot as plt
 from scipy.interpolate import interp1d
 from mbtrack2.collective_effects.wakefield import Impedance
 

collective_effects/wakefield.py

@@ -19,7 +19,9 @@ from scipy.integrate import trapz
 from scipy.constants import c
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 from mbtrack2.tracking.element import Element
-from mbtrack2.collective_effects.utilities import beam_spectrum, gaussian_bunch_spectrum, beam_loss_factor
+from mbtrack2.collective_effects.utilities import (beam_spectrum, gaussian_bunch_spectrum, 
+                                                   beam_loss_factor, spectral_density,
+                                                   effective_impedance)
 
 class ComplexData:
     """
@@ -193,6 +195,26 @@ class ComplexData:
 class WakeFunction(ComplexData):
     """
     Define a WakeFunction object based on a ComplexData object.
+    
+    Parameters
+    ----------
+    variable : array-like
+        Time domain structure of the wake function in [s].
+    function : array-like
+        Wake function values in [V/C].
+    wake_type : str, optinal
+        Type of the wake function: "long", "xdip", "xquad", ...
+    
+    Attributes
+    ----------
+    data : DataFrame
+    wake_type : str
+    
+    Methods
+    -------
+    from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
+        Compute a wake function from a wake potential file and load it to the 
+        WakeFunction object.
     """
 
     def __init__(self,
@@ -260,14 +282,22 @@ class WakeFunction(ComplexData):
     def wake_type(self, value):
         self._wake_type = value
         
-    def wakefunction_from_imp(self, imp_obj, nout=None, trim=False):
+    def from_wakepotential(self, file_name, bunch_length, bunch_charge, 
+                           freq_lim, nout=None, trim=False):
         """
-        Compute a wake function from an Impedance object and load it to 
-        the WakeFunction object.
+        Compute a wake function from a wake potential file and load it to the 
+        WakeFunction object.
     
         Parameters
         ----------
-        imp_obj : impedance object
+        file_name : str
+            Text file that contains wake potential data.
+        bunch_length : float
+            Spatial bunch length in [m].
+        bunch_charge : float
+            Bunch charge in [C].
+        freq_lim : float
+            The maximum frequency for calculating the impedance [Hz].
         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
@@ -279,68 +309,42 @@ class WakeFunction(ComplexData):
             If False, the original result is preserved.
             If a float is given, the pseudo wake function is trimmed from 
             time <= trim to 0. 
-        """
-        
-        Z0 = (imp_obj.data['real'] + imp_obj.data['imag']*1j)
-        Z = Z0[~np.isnan(Z0)]
-        
-        if imp_obj.impedance_type != "long":
-            Z = Z / 1j
-        
-        freq = Z.index
-        fs = ( freq[-1] - freq[0] ) / len(freq)
-        sampling = freq[1] - freq[0]
     
-        if nout is None:
-            nout = len(Z)
-            
-        time_array = sc.fft.fftfreq(nout, sampling)
-        Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs
-        
-        time = sc.fft.fftshift(time_array)
-        Wlong = sc.fft.fftshift(Wlong_raw)
-        
-        if trim is not False:
-            i_neg = np.where(time<trim)[0]
-            Wlong[i_neg] = 0
+        """
         
-        super().__init__(variable=time, function=Wlong)
-        self.data.index.name = "time [s]"        
+        imp = Impedance()
+        imp.from_wakepotential(file_name=file_name, bunch_length=bunch_length,
+                               bunch_charge=bunch_charge, freq_lim=freq_lim)
+        wf = imp.to_wakefunction(nout=nout, trim=trim)
+        self.__init__(variable=wf.data.index, function=wf.data["real"])
         
-    def long_wakefunction(self, Wp_filename, bunch_length, bunch_charge, 
-                          nout=None, freq_lim=200e9, trim=False):
+class Impedance(ComplexData):
     """
-        Compute wake function from wake potential and load it to the WakeFunction
-        object.
+    Define an Impedance object based on a ComplexData object.
     
     Parameters
     ----------
-        Wp_filename : str
-            Text file that contains wake potential data.
-        bunch_length : float
-            Spatial bunch length in [m].
-        bunch_charge : float
-            Bunch charge in [C].
-        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
-            input is padded with zeros. If nout < n, the inpu it cropped.
-            Note that, for any nout value, nout//2+1 input points are required.
-        freq_lim : float, optional
-            The maximum frequency for calculating the impedance [Hz].
+    variable : array-like
+        Frequency domain structure of the impedance in [Hz].
+    function : array-like
+        Impedance values in [Ohm].
+    impedance_type : str, optinal
+        Type of the impedance_type: "long", "xdip", "xquad", ...
     
-        """
-        
-        imp = Impedance()
-        imp.imp_from_wakepotential(Wp_filename=Wp_filename, 
-                                   freq_lim=freq_lim,
-                                   bunch_length=bunch_length, 
-                                   bunch_charge=bunch_charge)
-        self.wakefunction_from_imp(imp, nout=nout, trim=trim)
+    Attributes
+    ----------
+    data : DataFrame
+    impedance_type : str
     
-class Impedance(ComplexData):
-    """
-    Define an Impedance object based on a ComplexData object.
+    Methods
+    -------
+    from_wakepotential(file_name, bunch_length, bunch_charge, freq_lim)
+        Compute impedance from wake potential data and load it to the Impedance
+        object.
+    loss_factor(sigma)
+        Compute the loss factor or the kick factor for a Gaussian bunch.
+    to_wakefunction()
+        Return a WakeFunction object from the impedance data.
     """
 
     def __init__(self,
@@ -506,17 +510,23 @@ class Impedance(ComplexData):
 
         return kloss
     
-    def imp_from_wakepotential(self, Wp_filename, axis0='dist', 
-        axis0_scale=1e-3, axis1_scale=1e-12, freq_lim=200e9, 
-        bunch_length=1e-3, bunch_charge=1.44e-9):
+    def from_wakepotential(self, file_name, bunch_length, bunch_charge, 
+                           freq_lim, axis0='dist', axis0_scale=1e-3, 
+                           axis1_scale=1e-12):
         """
         Compute impedance from wake potential data and load it to the Impedance
         object.
 
         Parameters
         ----------
-        Wp_filename : str
+        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].
         axis0 : {'dist', 'time'}, optional
             Viariable of the data file's first column. Use 'dist' for spacial 
             distance. Otherwise, 'time' for temporal distance. 
@@ -525,16 +535,10 @@ class Impedance(ComplexData):
             or to second if axis0 = 'time'.
         axis1_scale : float, optional
             Scale of the wake potential (in second column) with respect to V/C.
-        freq_lim : float, optional
-            The maximum frequency for calculating the impedance [Hz].
-        bunch_length : float, optional
-            Electron bunch lenth [m]. 
-        bunch_charge : float, optional
-            Total bunch charge [C].
 
         """
         
-        s, wp0 = np.loadtxt(Wp_filename, unpack=True)
+        s, wp0 = np.loadtxt(file_name, unpack=True)
         if axis0 == 'dist':
             tau = s*axis0_scale/c
         elif axis0 == 'time':
@@ -565,6 +569,51 @@ class Impedance(ComplexData):
         super().__init__(variable=freq_trun, function=imp)
         self.data.index.name = "frequency [Hz]"    
         
+    def to_wakefunction(self, nout=None, trim=False):
+        """
+        Return a WakeFunction object from the impedance data.
+    
+        Parameters
+        ----------
+        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
+            input is padded with zeros. If nout < n, the inpu it cropped.
+            Note that, for any nout value, nout//2+1 input points are required.
+        trim : bool or float, optional
+            If True, the pseudo wake function is trimmed at time=0 and is forced 
+            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. 
+        """
+        
+        Z0 = (self.data['real'] + self.data['imag']*1j)
+        Z = Z0[~np.isnan(Z0)]
+        
+        if self.impedance_type != "long":
+            Z = Z / 1j
+        
+        freq = Z.index
+        fs = ( freq[-1] - freq[0] ) / len(freq)
+        sampling = freq[1] - freq[0]
+        
+        if nout is None:
+            nout = len(Z)
+            
+        time_array = sc.fft.fftfreq(nout, sampling)
+        Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs
+        
+        time = sc.fft.fftshift(time_array)
+        Wlong = sc.fft.fftshift(Wlong_raw)
+        
+        if trim is not False:
+            i_neg = np.where(time<trim)[0]
+            Wlong[i_neg] = 0
+                    
+        wf = WakeFunction(variable=time, function=Wlong)
+        return wf
+    
 class WakeField:
     """
     Defines a WakeField which corresponds to a single physical element which