From 231913a366d2f6384e207f532fcce549ebcbc83a Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 12 Aug 2022 16:16:34 +0200
Subject: [PATCH] Rework wakefield.py to improve WakeFunction class
Both Impedance and WakeFunction now inherit attributes from ComplexData (component_type, power_x, power_y, ...)
Add a components property to WakeField which lists all components present (W and Z)
Adapt other classes using Impedance and WakeFunction to the change.
---
mbtrack2/impedance/csr.py | 8 +-
mbtrack2/impedance/resistive_wall.py | 18 +-
mbtrack2/impedance/resonator.py | 12 +-
mbtrack2/impedance/tapers.py | 16 +-
mbtrack2/impedance/wakefield.py | 244 ++++++++++++++-------------
mbtrack2/utilities/misc.py | 10 +-
mbtrack2/utilities/read_impedance.py | 12 +-
7 files changed, 163 insertions(+), 157 deletions(-)
diff --git a/mbtrack2/impedance/csr.py b/mbtrack2/impedance/csr.py
index e7a2768..651e5be 100644
--- a/mbtrack2/impedance/csr.py
+++ b/mbtrack2/impedance/csr.py
@@ -47,8 +47,8 @@ class FreeSpaceCSR(WakeField):
Zl = self.LongitudinalImpedance(frequency)
# Wl = self.LongitudinalWakeFunction(time)
- Zlong = Impedance(variable = frequency, function = Zl, impedance_type='long')
- # Wlong = WakeFunction(variable = time, function = Wl, wake_type="long")
+ Zlong = Impedance(variable = frequency, function = Zl, component_type='long')
+ # Wlong = WakeFunction(variable = time, function = Wl, component_type="long")
super().append_to_model(Zlong)
# super().append_to_model(Wlong)
@@ -126,8 +126,8 @@ class ParallelPlatesCSR(WakeField):
Zl = self.LongitudinalImpedance(frequency)
# Wl = self.LongitudinalWakeFunction(time)
- Zlong = Impedance(variable = frequency, function = Zl, impedance_type='long')
- # Wlong = WakeFunction(variable = time, function = Wl, wake_type="long")
+ Zlong = Impedance(variable = frequency, function = Zl, component_type='long')
+ # Wlong = WakeFunction(variable = time, function = Wl, component_type="long")
super().append_to_model(Zlong)
# super().append_to_model(Wlong)
diff --git a/mbtrack2/impedance/resistive_wall.py b/mbtrack2/impedance/resistive_wall.py
index c7f13a8..8434db6 100644
--- a/mbtrack2/impedance/resistive_wall.py
+++ b/mbtrack2/impedance/resistive_wall.py
@@ -94,12 +94,12 @@ class CircularResistiveWall(WakeField):
Wl = self.LongitudinalWakeFunction(time, exact, atol)
Wt = self.TransverseWakeFunction(time, exact)
- Zlong = Impedance(variable = frequency, function = Z1, impedance_type='long')
- Zxdip = Impedance(variable = frequency, function = Z2, impedance_type='xdip')
- Zydip = Impedance(variable = frequency, function = Z2, impedance_type='ydip')
- Wlong = WakeFunction(variable = time, function = Wl, wake_type="long")
- Wxdip = WakeFunction(variable = time, function = Wt, wake_type="xdip")
- Wydip = WakeFunction(variable = time, function = Wt, wake_type="ydip")
+ Zlong = Impedance(variable = frequency, function = Z1, component_type='long')
+ Zxdip = Impedance(variable = frequency, function = Z2, component_type='xdip')
+ Zydip = Impedance(variable = frequency, function = Z2, component_type='ydip')
+ Wlong = WakeFunction(variable = time, function = Wl, component_type="long")
+ Wxdip = WakeFunction(variable = time, function = Wt, component_type="xdip")
+ Wydip = WakeFunction(variable = time, function = Wt, component_type="ydip")
super().append_to_model(Zlong)
super().append_to_model(Zxdip)
@@ -280,9 +280,9 @@ class Coating(WakeField):
Zl = self.LongitudinalImpedance(frequency, approx)
Zt = self.TransverseImpedance(frequency, approx)
- Zlong = Impedance(variable = frequency, function = Zl, impedance_type='long')
- Zxdip = Impedance(variable = frequency, function = Zt, impedance_type='xdip')
- Zydip = Impedance(variable = frequency, function = Zt, impedance_type='ydip')
+ Zlong = Impedance(variable = frequency, function = Zl, component_type='long')
+ Zxdip = Impedance(variable = frequency, function = Zt, component_type='xdip')
+ Zydip = Impedance(variable = frequency, function = Zt, component_type='ydip')
super().append_to_model(Zlong)
super().append_to_model(Zxdip)
diff --git a/mbtrack2/impedance/resonator.py b/mbtrack2/impedance/resonator.py
index d514e52..503d106 100644
--- a/mbtrack2/impedance/resonator.py
+++ b/mbtrack2/impedance/resonator.py
@@ -59,21 +59,21 @@ class Resonator(WakeField):
if dim == "long":
Zlong = Impedance(variable=frequency,
function=self.long_impedance(frequency),
- impedance_type="long")
+ component_type="long")
super().append_to_model(Zlong)
Wlong = WakeFunction(variable=time,
function=self.long_wake_function(time, atol),
- wake_type="long")
+ component_type="long")
super().append_to_model(Wlong)
elif dim == "x" or dim == "y":
Zdip = Impedance(variable=frequency,
function=self.transverse_impedance(frequency),
- impedance_type=dim + "dip")
+ component_type=dim + "dip")
super().append_to_model(Zdip)
Wdip = WakeFunction(variable=time,
function=self.transverse_wake_function(time),
- wake_type=dim + "dip")
+ component_type=dim + "dip")
super().append_to_model(Wdip)
else:
raise ValueError("Plane must be: long, x or y")
@@ -137,7 +137,7 @@ class PureInductive(WakeField):
freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
imp = Impedance(variable=freq,
function=self.long_impedance(freq),
- impedance_type="long")
+ component_type="long")
super().append_to_model(imp)
wf = imp.to_wakefunction(nout=nout, trim=trim)
@@ -173,7 +173,7 @@ class PureResistive(WakeField):
freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
imp = Impedance(variable=freq,
function=self.long_impedance(freq),
- impedance_type="long")
+ component_type="long")
super().append_to_model(imp)
wf = imp.to_wakefunction(nout=nout, trim=trim)
diff --git a/mbtrack2/impedance/tapers.py b/mbtrack2/impedance/tapers.py
index 838e275..e9d575b 100644
--- a/mbtrack2/impedance/tapers.py
+++ b/mbtrack2/impedance/tapers.py
@@ -36,11 +36,11 @@ class StupakovRectangularTaper(WakeField):
self.m_max = m_max
self.n_points = n_points
- Zlong = Impedance(variable = frequency, function = self.long(), impedance_type='long')
- Zxdip = Impedance(variable = frequency, function = self.xdip(), impedance_type='xdip')
- Zydip = Impedance(variable = frequency, function = self.ydip(), impedance_type='ydip')
- Zxquad = Impedance(variable = frequency, function = -1*self.quad(), impedance_type='xquad')
- Zyquad = Impedance(variable = frequency, function = self.quad(), impedance_type='yquad')
+ Zlong = Impedance(variable = frequency, function = self.long(), component_type='long')
+ Zxdip = Impedance(variable = frequency, function = self.xdip(), component_type='xdip')
+ Zydip = Impedance(variable = frequency, function = self.ydip(), component_type='ydip')
+ Zxquad = Impedance(variable = frequency, function = -1*self.quad(), component_type='xquad')
+ Zyquad = Impedance(variable = frequency, function = self.quad(), component_type='yquad')
super().append_to_model(Zlong)
super().append_to_model(Zxdip)
@@ -160,9 +160,9 @@ class StupakovCircularTaper(WakeField):
self.m_max = m_max
self.n_points = n_points
- Zlong = Impedance(variable = frequency, function = self.long(), impedance_type='long')
- Zxdip = Impedance(variable = frequency, function = self.dip(), impedance_type='xdip')
- Zydip = Impedance(variable = frequency, function = self.dip(), impedance_type='ydip')
+ Zlong = Impedance(variable = frequency, function = self.long(), component_type='long')
+ Zxdip = Impedance(variable = frequency, function = self.dip(), component_type='xdip')
+ Zydip = Impedance(variable = frequency, function = self.dip(), component_type='ydip')
super().append_to_model(Zlong)
super().append_to_model(Zxdip)
diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py
index 6a3fd12..f21b317 100644
--- a/mbtrack2/impedance/wakefield.py
+++ b/mbtrack2/impedance/wakefield.py
@@ -23,7 +23,7 @@ class ComplexData:
variable : list, numpy array
contains the function variable values
- function : list or numpy array of complex numbers
+ function : list or numpy array of comp lex numbers
contains the values taken by the complex function
"""
@@ -182,6 +182,71 @@ class ComplexData:
y = self.data["imag"], kind=interp_kind)
return real_func(values) + 1j*imag_func(values)
+
+ def initialize_coefficient(self):
+ """
+ Define the impedance coefficients and the plane of the impedance that
+ are presents in attributes of the class.
+ """
+ table = self.name_and_coefficients_table()
+
+ try:
+ component_coefficients = table[self.component_type].to_dict()
+ except KeyError:
+ print('Component type {} does not exist'.format(self.component_type))
+ raise
+
+ self.a = component_coefficients["a"]
+ self.b = component_coefficients["b"]
+ self.c = component_coefficients["c"]
+ self.d = component_coefficients["d"]
+ self.plane = component_coefficients["plane"]
+
+ def name_and_coefficients_table(self):
+ """
+ Return a table associating the human readbale names of an impedance
+ component and its associated coefficients and plane.
+ """
+
+ component_dict = {
+ 'long': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'z'},
+ 'xcst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
+ 'ycst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
+ 'xdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
+ 'ydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'y'},
+ 'xydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'x'},
+ 'yxdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
+ 'xquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'x'},
+ 'yquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'y'},
+ 'xyquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'x'},
+ 'yxquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'y'},
+ }
+
+ return pd.DataFrame(component_dict)
+
+ @property
+ def power_x(self):
+ power_x = self.a/2 + self.c/2.
+ if self.plane == 'x':
+ power_x += 1./2.
+ return power_x
+
+ @property
+ def power_y(self):
+ power_y = self.b/2. + self.d/2.
+ if self.plane == 'y':
+ power_y += 1./2.
+ return power_y
+
+ @property
+ def component_type(self):
+ return self._component_type
+
+ @component_type.setter
+ def component_type(self, value):
+ self._component_type = value
+ self.initialize_coefficient()
+
class WakeFunction(ComplexData):
"""
Define a WakeFunction object based on a ComplexData object.
@@ -192,7 +257,7 @@ class WakeFunction(ComplexData):
Time domain structure of the wake function in [s].
function : array-like
Wake function values in [V/C].
- wake_type : str, optinal
+ component_type : str, optinal
Type of the wake function: "long", "xdip", "xquad", ...
Attributes
@@ -209,10 +274,11 @@ class WakeFunction(ComplexData):
def __init__(self,
variable=np.array([-1e15, 1e15]),
- function=np.array([0, 0]), wake_type='long'):
+ function=np.array([0, 0]), component_type='long'):
super().__init__(variable, function)
- self._wake_type = wake_type
+ self._component_type = component_type
self.data.index.name = "time [s]"
+ self.initialize_coefficient()
def add(self, structure_to_add, method='zero'):
"""
@@ -224,7 +290,7 @@ class WakeFunction(ComplexData):
result = super().add(structure_to_add, method=method,
index_name="time [s]")
elif isinstance(structure_to_add, WakeFunction):
- if (self.wake_type == structure_to_add.wake_type):
+ if (self.component_type == structure_to_add.component_type):
result = super().add(structure_to_add, method=method,
index_name="time [s]")
else:
@@ -237,7 +303,7 @@ class WakeFunction(ComplexData):
wake_to_return = WakeFunction(
result.data.index,
result.data.real.values,
- self.wake_type)
+ self.component_type)
return wake_to_return
def __radd__(self, structure_to_add):
@@ -255,7 +321,7 @@ class WakeFunction(ComplexData):
wake_to_return = WakeFunction(
result.data.index,
result.data.real.values,
- self.wake_type)
+ self.component_type)
return wake_to_return
def __mul__(self, factor):
@@ -263,17 +329,9 @@ class WakeFunction(ComplexData):
def __rmul__(self, factor):
return self.multiply(factor)
-
- @property
- def wake_type(self):
- return self._wake_type
-
- @wake_type.setter
- def wake_type(self, value):
- self._wake_type = value
def from_wakepotential(self, file_name, bunch_length, bunch_charge,
- freq_lim, wake_type="long", divide_by=None,
+ freq_lim, component_type="long", divide_by=None,
nout=None, trim=False, axis0='dist',
axis0_scale=1e-3, axis1_scale=1e-12):
"""
@@ -290,7 +348,7 @@ class WakeFunction(ComplexData):
Bunch charge in [C].
freq_lim : float
The maximum frequency for calculating the impedance [Hz].
- wake_type : str, optional
+ component_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
@@ -320,12 +378,12 @@ class WakeFunction(ComplexData):
imp = Impedance()
imp.from_wakepotential(file_name=file_name, bunch_length=bunch_length,
bunch_charge=bunch_charge, freq_lim=freq_lim,
- impedance_type=wake_type, divide_by=divide_by,
+ component_type=component_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"],
- wake_type=wake_type)
+ component_type=component_type)
class Impedance(ComplexData):
"""
@@ -337,8 +395,8 @@ class Impedance(ComplexData):
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", ...
+ component_type : str, optinal
+ Type of the impedance: "long", "xdip", "xquad", ...
Attributes
----------
@@ -358,53 +416,11 @@ class Impedance(ComplexData):
def __init__(self,
variable=np.array([-1e15, 1e15]),
- function=np.array([0, 0]), impedance_type='long'):
+ function=np.array([0, 0]), component_type='long'):
super().__init__(variable, function)
- self._impedance_type = impedance_type
+ self._component_type = component_type
self.data.index.name = 'frequency [Hz]'
- self.initialize_impedance_coefficient()
-
-
- def initialize_impedance_coefficient(self):
- """
- Define the impedance coefficients and the plane of the impedance that
- are presents in attributes of the class.
- """
- table = self.impedance_name_and_coefficients_table()
-
- try:
- 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"]
- self.c = component_coefficients["c"]
- self.d = component_coefficients["d"]
- self.plane = component_coefficients["plane"]
-
- def impedance_name_and_coefficients_table(self):
- """
- Return a table associating the human readbale names of an impedance
- component and its associated coefficients and plane.
- """
-
- component_dict = {
- 'long': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'z'},
- 'xcst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
- 'ycst': {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
- 'xdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'x'},
- 'ydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'y'},
- 'xydip': {'a': 0, 'b': 1, 'c': 0, 'd': 0, 'plane': 'x'},
- 'yxdip': {'a': 1, 'b': 0, 'c': 0, 'd': 0, 'plane': 'y'},
- 'xquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'x'},
- 'yquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'y'},
- 'xyquad': {'a': 0, 'b': 0, 'c': 0, 'd': 1, 'plane': 'x'},
- 'yxquad': {'a': 0, 'b': 0, 'c': 1, 'd': 0, 'plane': 'y'},
- }
-
- return pd.DataFrame(component_dict)
+ self.initialize_coefficient()
def add(self, structure_to_add, beta_x=1, beta_y=1, method='zero'):
"""
@@ -417,7 +433,7 @@ class Impedance(ComplexData):
result = super().add(structure_to_add, method=method,
index_name="frequency [Hz]")
elif isinstance(structure_to_add, Impedance):
- if (self.impedance_type == structure_to_add.impedance_type):
+ if (self.component_type == structure_to_add.component_type):
weight = (beta_x ** self.power_x) * (beta_y ** self.power_y)
result = super().add(structure_to_add * weight, method=method,
index_name="frequency [Hz]")
@@ -431,7 +447,7 @@ class Impedance(ComplexData):
impedance_to_return = Impedance(
result.data.index,
result.data.real.values + 1j*result.data.imag.values,
- self.impedance_type)
+ self.component_type)
return impedance_to_return
def __radd__(self, structure_to_add):
@@ -449,7 +465,7 @@ class Impedance(ComplexData):
impedance_to_return = Impedance(
result.data.index,
result.data.real.values + 1j*result.data.imag.values,
- self.impedance_type)
+ self.component_type)
return impedance_to_return
def __mul__(self, factor):
@@ -457,29 +473,6 @@ class Impedance(ComplexData):
def __rmul__(self, factor):
return self.multiply(factor)
-
- @property
- def impedance_type(self):
- return self._impedance_type
-
- @impedance_type.setter
- def impedance_type(self, value):
- self._impedance_type = value
- self.initialize_impedance_coefficient()
-
- @property
- def power_x(self):
- power_x = self.a/2 + self.c/2.
- if self.plane == 'x':
- power_x += 1./2.
- return power_x
-
- @property
- def power_y(self):
- power_y = self.b/2. + self.d/2.
- if self.plane == 'y':
- power_y += 1./2.
- return power_y
def loss_factor(self, sigma):
"""
@@ -509,13 +502,13 @@ class Impedance(ComplexData):
from mbtrack2.utilities import spectral_density
sd = spectral_density(frequency, sigma, m=0)
- if(self.impedance_type == "long"):
+ if(self.component_type == "long"):
kloss = trapz(x = frequency,
y = 2*self.data["real"][positive_index]*sd)
- elif(self.impedance_type == "xdip" or self.impedance_type == "ydip"):
+ elif(self.component_type == "xdip" or self.component_type == "ydip"):
kloss = trapz(x = frequency,
y = 2*self.data["imag"][positive_index]*sd)
- elif(self.impedance_type == "xquad" or self.impedance_type == "yquad"):
+ elif(self.component_type == "xquad" or self.component_type == "yquad"):
kloss = trapz(x = frequency,
y = 2*self.data["imag"][positive_index]*sd)
else:
@@ -524,7 +517,7 @@ class Impedance(ComplexData):
return kloss
def from_wakepotential(self, file_name, bunch_length, bunch_charge,
- freq_lim, impedance_type="long", divide_by=None,
+ freq_lim, component_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
@@ -540,7 +533,7 @@ class Impedance(ComplexData):
Total bunch charge [C].
freq_lim : float
The maximum frequency for calculating the impedance [Hz].
- impedance_type : str, optional
+ component_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
@@ -584,15 +577,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
- if impedance_type == "long":
+ if component_type == "long":
imp = dft_wake_trun/dft_rho_trun*-1*bunch_charge
- elif (impedance_type == "xdip") or (impedance_type == "ydip"):
+ elif (component_type == "xdip") or (component_type == "ydip"):
imp = dft_wake_trun/dft_rho_trun*-1j*bunch_charge
else:
- raise NotImplementedError(impedance_type + " is not correct.")
+ raise NotImplementedError(component_type + " is not correct.")
self.__init__(variable=freq_trun, function=imp,
- impedance_type=impedance_type)
+ component_type=component_type)
def to_wakefunction(self, nout=None, trim=False):
"""
@@ -617,7 +610,7 @@ class Impedance(ComplexData):
Z0 = (self.data['real'] + self.data['imag']*1j)
Z = Z0[~np.isnan(Z0)]
- if self.impedance_type != "long":
+ if self.component_type != "long":
Z = Z / 1j
freq = Z.index
@@ -640,7 +633,7 @@ class Impedance(ComplexData):
Wlong[i_neg] = 0
wf = WakeFunction(variable=time, function=Wlong,
- wake_type=self.impedance_type)
+ component_type=self.component_type)
return wf
class WakeField:
@@ -658,7 +651,11 @@ class WakeField:
Attributes
----------
impedance_components : array of str
- Impedance component names present for this element.
+ Impedance components present for this element.
+ wake_components : array of str
+ WakeFunction components present for this element.
+ components : array of str
+ Impedance or WakeFunction components present for this element.
Methods
-------
@@ -686,14 +683,14 @@ class WakeField:
"""
list_of_attributes = dir(self)
if isinstance(structure_to_add, Impedance):
- attribute_name = "Z" + structure_to_add.impedance_type
+ attribute_name = "Z" + structure_to_add.component_type
if attribute_name in list_of_attributes:
raise ValueError("There is already a component of the type "
"{} in this element.".format(attribute_name))
else:
self.__setattr__(attribute_name, structure_to_add)
elif isinstance(structure_to_add, WakeFunction):
- attribute_name = "W" + structure_to_add.wake_type
+ attribute_name = "W" + structure_to_add.component_type
if attribute_name in list_of_attributes:
raise ValueError("There is already a component of the type "
"{} in this element.".format(attribute_name))
@@ -730,6 +727,15 @@ class WakeField:
valid = ["Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad"]
return np.array([comp for comp in dir(self) if comp in valid])
+ @property
+ def components(self):
+ """
+ Return an array of the component names for the element.
+ """
+ valid = ["Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad",
+ "Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad"]
+ return np.array([comp for comp in dir(self) if comp in valid])
+
@staticmethod
def add_wakefields(wake1, beta1, wake2, beta2):
"""
@@ -753,22 +759,22 @@ class WakeField:
"""
wake_sum = deepcopy(wake1)
- for component_name1 in wake1.impedance_components:
- impedance1 = getattr(wake_sum, component_name1)
- weight1 = ((beta1[0] ** impedance1.power_x) *
- (beta1[1] ** impedance1.power_y))
- setattr(wake_sum, component_name1, weight1*impedance1)
+ for component_name1 in wake1.components:
+ comp1 = getattr(wake_sum, component_name1)
+ weight1 = ((beta1[0] ** comp1.power_x) *
+ (beta1[1] ** comp1.power_y))
+ setattr(wake_sum, component_name1, weight1*comp1)
- for component_name2 in wake2.impedance_components:
- impedance2 = getattr(wake2, component_name2)
- weight2 = ((beta2[0] ** impedance2.power_x) *
- (beta2[1] ** impedance2.power_y))
+ for component_name2 in wake2.components:
+ comp2 = getattr(wake2, component_name2)
+ weight2 = ((beta2[0] ** comp2.power_x) *
+ (beta2[1] ** comp2.power_y))
try:
- impedance1 = getattr(wake_sum, component_name2)
- setattr(wake_sum, component_name2, impedance1 +
- weight2*impedance2)
+ comp1 = getattr(wake_sum, component_name2)
+ setattr(wake_sum, component_name2, comp1 +
+ weight2*comp2)
except AttributeError:
- setattr(wake_sum, component_name2, weight2*impedance2)
+ setattr(wake_sum, component_name2, weight2*comp2)
return wake_sum
diff --git a/mbtrack2/utilities/misc.py b/mbtrack2/utilities/misc.py
index a099c2a..93332b8 100644
--- a/mbtrack2/utilities/misc.py
+++ b/mbtrack2/utilities/misc.py
@@ -68,19 +68,19 @@ def effective_impedance(ring, imp, m, mu, sigma, M, tuneS, xi=None,
p = np.arange(pmin,pmax+1)
- if imp.impedance_type == "long":
+ if imp.component_type == "long":
fp = ring.f0*(p*M + mu + m*tuneS)
fp = fp[np.nonzero(fp)] # Avoid division by 0
num = np.sum( imp(fp) * h(fp) / (fp*2*np.pi) )
den = np.sum( h(fp) )
Zeff = num/den
- elif imp.impedance_type == "xdip" or imp.impedance_type == "ydip":
- if imp.impedance_type == "xdip":
+ elif imp.component_type == "xdip" or imp.component_type == "ydip":
+ if imp.component_type == "xdip":
tuneXY = ring.tune[0]
if xi is None :
xi = ring.chro[0]
- elif imp.impedance_type == "ydip":
+ elif imp.component_type == "ydip":
tuneXY = ring.tune[1]
if xi is None:
xi = ring.chro[1]
@@ -214,7 +214,7 @@ def double_sided_impedance(impedance):
negative_index = impedance.data.index*-1
negative_data = impedance.data.set_index(negative_index)
- imp_type = impedance.impedance_type
+ imp_type = impedance.component_type
if imp_type == "long":
negative_data["imag"] = -1*negative_data["imag"]
diff --git a/mbtrack2/utilities/read_impedance.py b/mbtrack2/utilities/read_impedance.py
index 2ea65a8..9d0ac30 100644
--- a/mbtrack2/utilities/read_impedance.py
+++ b/mbtrack2/utilities/read_impedance.py
@@ -10,7 +10,7 @@ from pathlib import Path
from scipy.constants import c
from mbtrack2.impedance.wakefield import Impedance, WakeFunction, WakeField
-def read_CST(file, impedance_type='long', divide_by=None):
+def read_CST(file, component_type='long', divide_by=None):
"""
Read CST file format into an Impedance object.
@@ -18,7 +18,7 @@ def read_CST(file, impedance_type='long', divide_by=None):
----------
file : str
Path to the file to read.
- impedance_type : str, optional
+ component_type : str, optional
Type of the Impedance object.
divide_by : float, optional
Divide the impedance by a value. Mainly used to normalize transverse
@@ -35,12 +35,12 @@ 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":
+ if component_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"],
- impedance_type=impedance_type)
+ component_type=component_type)
return result
def read_IW2D(file, file_type='Zlong'):
@@ -67,7 +67,7 @@ def read_IW2D(file, file_type='Zlong'):
df = df[df["Imaginary"].notna()]
result = Impedance(variable = df.index,
function = df["Real"] + 1j*df["Imaginary"],
- impedance_type=file_type[1:])
+ component_type=file_type[1:])
elif file_type[0] == "W":
df = pd.read_csv(file, delim_whitespace=True, header = None,
names = ["Distance","Wake"], skiprows=1)
@@ -82,7 +82,7 @@ def read_IW2D(file, file_type='Zlong'):
# df["Wake"] = df["Wake"]*-1
result = WakeFunction(variable = df.index,
function = df["Wake"],
- wake_type=file_type[1:])
+ component_type=file_type[1:])
else:
raise ValueError("file_type should begin by Z or W.")
return result
--
GitLab