diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py
index a8ff83234c9585f87bed18d0b2818413dd083fd1..6a3fd120fbcb556535b4722bd376349713d2d16f 100644
--- a/mbtrack2/impedance/wakefield.py
+++ b/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:
diff --git a/mbtrack2/utilities/read_impedance.py b/mbtrack2/utilities/read_impedance.py
index 5a3f0a501a433e37996ad6d16dc0985f316aa0f0..2ea65a8b929803310132bd28ef77e2624829394c 100644
--- a/mbtrack2/utilities/read_impedance.py
+++ b/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)