Skip to content
Snippets Groups Projects
Commit 231913a3 authored by Alexis GAMELIN's avatar Alexis GAMELIN
Browse files

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.
parent 0b61c7ae
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
mbtrack2/impedance/csr.py
+ 4
− 4
View file @ 231913a3
......@@ -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)
......
mbtrack2/impedance/resistive_wall.py
+ 9
− 9
View file @ 231913a3
......@@ -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)
......
mbtrack2/impedance/resonator.py
+ 6
− 6
View file @ 231913a3
......@@ -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)
......
mbtrack2/impedance/tapers.py
+ 8
− 8
View file @ 231913a3
......@@ -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)
......
mbtrack2/impedance/wakefield.py
+ 125
− 119
View file @ 231913a3
......@@ -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
......
mbtrack2/utilities/misc.py
+ 5
− 5
View file @ 231913a3
......@@ -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"]
......
mbtrack2/utilities/read_impedance.py
+ 6
− 6
View file @ 231913a3
......@@ -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
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

The place to be for your source code