diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index ebc7b09debed36521c310dd2884f296260d4cce7..bc3da3c0f8f7b3bb2794e580668c1efdc1783239 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -11,26 +11,28 @@ from copy import deepcopy
from scipy.integrate import trapz
from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
-from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance,
+from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance,
double_sided_impedance)
-from mbtrack2.utilities.spectrum import (beam_spectrum,
- gaussian_bunch_spectrum,
+from mbtrack2.utilities.spectrum import (beam_spectrum,
+ gaussian_bunch_spectrum,
spectral_density)
+from mbtrack2.impedance.wakefield import WakeField, WakeFunction
+
class ImpedanceModel():
"""
Define the impedance model of the machine.
-
- The model must be completed with successive add(...) and add_global(...)
+
+ The model must be completed with successive add(...) and add_global(...)
calls, then compute_sum() must be run.
-
- The transverse impedance and wake functions are beta weighted and divided
+
+ The transverse impedance and wake functions are beta weighted and divided
by the beta at the tracking location (ring.optics.local_beta).
-
+
Parameters
----------
ring : Synchrotron object
-
+
Attributes
----------
wakefields : list of WakeField objects
@@ -44,12 +46,12 @@ class ImpedanceModel():
sum_"name" : WakeField
Sum of the "name" Wakefield weigthed by beta functions.
sum_names : array
- Names of attributes where the WakeFields are summed by name.
+ Names of attributes where the WakeFields are summed by name.
globals : list of WakeField objects
Globals WakeFields in the model.
globals_names : list of str
Names of the global WakeFields objects.
-
+
Methods
-------
add(wakefield, positions, name)
@@ -70,7 +72,7 @@ class ImpedanceModel():
load(file)
Load impedance model from file.
"""
-
+
def __init__(self, ring):
self.ring = ring
self.optics = self.ring.optics
@@ -84,7 +86,7 @@ class ImpedanceModel():
def add(self, wakefield, positions, name=None):
"""
Add the same WakeField object at different locations to the model.
-
+
Parameters
----------
wakefield : WakeField
@@ -92,7 +94,7 @@ class ImpedanceModel():
positions : array, float or int
Array of longitudinal positions where the elements are loacted.
name : str, optional
- Name of the element type. If None, the name of the WakeField object
+ Name of the element type. If None, the name of the WakeField object
is used. The default is None.
Returns
@@ -110,22 +112,22 @@ class ImpedanceModel():
self.names.append(name)
else:
raise ValueError("This name is already taken.")
-
+
def add_global(self, wakefield, name=None):
"""
Add a "global" WakeField object which will added to the sum WakeField
but not weighted by beta functions.
-
- To use with "distributed" elements, for example a resistive wall
- wakefield computed from an effective radius (so which has already been
- weighted by beta functions).
+
+ To use with "distributed" elements, for example a resistive wall
+ wakefield computed from an effective radius (so which has already been
+ weighted by beta functions).
Parameters
----------
wakefield : WakeField
WakeField object to add to the model.
name : str, optional
- Name of the element type. If None, the name of the WakeField object
+ Name of the element type. If None, the name of the WakeField object
is used. The default is None.
Returns
@@ -143,7 +145,7 @@ class ImpedanceModel():
else:
raise ValueError("This name is already taken.")
# setattr(self, name, wakefield)
-
+
def sum_beta(self, wake, beta):
"""
Weight a WakeField object by an array of beta functions.
@@ -165,8 +167,8 @@ class ImpedanceModel():
local_beta = self.ring.optics.local_beta
for component_name in wake_sum.components:
comp = getattr(wake_sum, component_name)
- weight = ((beta[0,:] ** comp.power_x) *
- (beta[1,:] ** comp.power_y))
+ weight = ((beta[0, :] ** comp.power_x) *
+ (beta[1, :] ** comp.power_y))
if comp.plane == "x":
weight = weight.sum() / local_beta[0]
if comp.plane == "y":
@@ -175,7 +177,7 @@ class ImpedanceModel():
weight = weight.sum()
setattr(wake_sum, component_name, weight*comp)
return wake_sum
-
+
def compute_sum_names(self):
"""
Compute the weighted WakeField for each WakeField object type.
@@ -188,19 +190,19 @@ class ImpedanceModel():
wake_sum.name = attribute_name
setattr(self, attribute_name, wake_sum)
self.sum_names.append(attribute_name)
-
+
def compute_sum(self):
"""
Compute the sum of all weighted WakeField into self.sum.
"""
self.compute_sum_names()
for i, name in enumerate(self.sum_names):
- if i==0:
+ if i == 0:
self.sum = deepcopy(getattr(self, name))
self.sum.name = "sum"
else:
wake2 = getattr(self, name)
- for component_name2 in wake2.components:
+ for component_name2 in wake2.components:
comp2 = getattr(wake2, component_name2)
try:
comp1 = getattr(self.sum, component_name2)
@@ -215,7 +217,7 @@ class ImpedanceModel():
self.sum = deepcopy(wake2)
self.sum.name = "sum"
else:
- for component_name2 in wake2.components:
+ for component_name2 in wake2.components:
comp2 = getattr(wake2, component_name2)
try:
comp1 = getattr(self.sum, component_name2)
@@ -223,7 +225,64 @@ class ImpedanceModel():
except AttributeError:
setattr(self.sum, component_name2, comp2)
- def plot_area(self, Z_type="Zlong", component="real", sigma=None,
+ def group_attributes(self, string_in_name, names_to_group=None, property_list=['Zlong']):
+ """Groups attributes in the ImpedanceModel based on a given string pattern.
+ Args:
+ string_in_name (str): The string pattern used to match attribute names for grouping. If names_to_group is given, this is a name of a new attribute instead.
+ names_to_group (list, optional): List of attribute names to be explicitly grouped.
+ If not provided, attributes matching the string pattern will be automatically grouped.
+ Defaults to None.
+ property_list (list, optional): List of property names to be grouped for each attribute.
+ Defaults to ['Zlong'].
+ Returns:
+ int: Returns 0 indicating the successful grouping of attributes.
+ Notes:
+ - This method groups attributes in the ImpedanceModel class based on a given string pattern.
+ - If names_to_group is not provided, it automatically selects attributes from the existing sum_names that match the given string pattern.
+ - A new WakeField instance is created and assigned to the attribute named 'string_in_name'.
+ - The specified properties from the first attribute in names_to_group are appended to the new WakeField instance.
+ - The values of the specified properties from the remaining attributes in names_to_group are added to the corresponding properties in the new WakeField instance.
+ - The names_to_group attributes are removed from the ImpedanceModel class.
+ - The new grouped attribute is added to sum_names, and the first attribute in names_to_group is removed from sum_names."""
+ attrs = self.sum_names
+ if names_to_group == None:
+ names_to_group = []
+ for attr in attrs:
+ if string_in_name in attr:
+ names_to_group.append(attr)
+ setattr(self, string_in_name, WakeField())
+ for prop in property_list:
+ getattr(self, string_in_name).append_to_model(
+ getattr(getattr(self, names_to_group[0]), prop))
+ for attr in names_to_group[1:]:
+ for prop in property_list:
+ old_values = getattr(getattr(self, string_in_name), prop)
+ new_values = getattr(getattr(self, attr), prop)
+ setattr(getattr(self, string_in_name),
+ prop, old_values+new_values)
+ self.sum_names.remove(attr)
+ delattr(self, attr)
+ self.sum_names.append(string_in_name)
+ self.sum_names.remove(names_to_group[0])
+ delattr(self, names_to_group[0])
+ return 0
+
+ def rename_attribute(self, old_name, new_name):
+ """Renames an attribute in the ImpedanceModel.
+ Args:
+ old_name (str): The current name of the attribute to be renamed.
+ new_name (str): The new name for the attribute.
+ Raises:
+ KeyError: If the old_name doesn't exist as an attribute in the ImpedanceModel.
+ Notes:
+ - This method renames an attribute in the ImpedanceModel class.
+ - The attribute with the old_name is removed from the class's dictionary (__dict__) using the pop() method.
+ - The attribute is then added back to the class's dictionary with the new_name using the __dict__ attribute.
+ - If the old_name doesn't exist as an attribute in the ImpedanceModel, a KeyError is raised.
+ """
+ self.__dict__[new_name] = self.__dict__.pop(old_name)
+
+ def plot_area(self, Z_type="Zlong", component="real", sigma=None,
attr_list=None, zoom=False):
"""
Plot the contributions of different kind of WakeFields.
@@ -246,46 +305,47 @@ class ImpedanceModel():
if attr_list is None:
attr_list = self.sum_names
-
# manage legend
- Ztype_dict = {"Zlong":0, "Zxdip":1, "Zydip":2, "Zxquad":3, "Zyquad":4}
+ Ztype_dict = {"Zlong": 0, "Zxdip": 1,
+ "Zydip": 2, "Zxquad": 3, "Zyquad": 4}
scale = [1e-3, 1e-6, 1e-6, 1e-6, 1e-6]
- label_list = [r"$Z_\mathrm{long} \; (\mathrm{k}\Omega)$",
- r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
- r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
- r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^\mathrm{Quad} \; (\mathrm{M}\Omega/\mathrm{m})$",
- r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^\mathrm{Quad} \; (\mathrm{M}\Omega/\mathrm{m})$"]
+ label_list = [r"$Z_\mathrm{long} \; (\mathrm{k}\Omega)$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^\mathrm{Quad} \; (\mathrm{M}\Omega/\mathrm{m})$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^\mathrm{Quad} \; (\mathrm{M}\Omega/\mathrm{m})$"]
leg = Ztype_dict[Z_type]
-
- # sort plot by decresing area size
+
+ # sort plot by decresing area size
area = np.zeros((len(attr_list),))
for index, attr in enumerate(attr_list):
try:
sum_imp = getattr(getattr(self, attr), Z_type)
- area[index] = trapz(sum_imp.data[component], sum_imp.data.index)
+ area[index] = trapz(
+ sum_imp.data[component], sum_imp.data.index)
except AttributeError:
pass
sorted_index = area.argsort()[::-1]
-
+
# Init fig
fig = plt.figure()
ax = fig.gca()
zero_impedance = getattr(self.sum, Z_type)*0
total_imp = 0
legend = []
-
+
if sigma is not None:
legend.append("Spectral density for sigma = " + str(sigma) + " s")
-
+
# Main plot
colorblind = colormaps['tab10'].colors
- for index in sorted_index:
+ for index in sorted_index:
attr = attr_list[index]
# Set all impedances with common indexes using + zero_impedance
try:
sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
- ax.fill_between(sum_imp.data.index*1e-9, total_imp,
- total_imp + sum_imp.data[component]*scale[leg], edgecolor=colorblind[index%10], color=colorblind[index%10])
+ ax.fill_between(sum_imp.data.index*1e-9, total_imp,
+ total_imp + sum_imp.data[component]*scale[leg], edgecolor=colorblind[index % 10], color=colorblind[index % 10])
total_imp += sum_imp.data[component]*scale[leg]
if attr[:4] == "sum_":
legend.append(attr[4:])
@@ -293,41 +353,42 @@ class ImpedanceModel():
legend.append(attr)
except AttributeError:
pass
-
+
if sigma is not None:
spect = spectral_density(zero_impedance.data.index, sigma)
spect = spect/spect.max()*total_imp.max()
ax.plot(sum_imp.data.index*1e-9, spect, 'r', linewidth=2.5)
-
- ax.legend(legend, loc="upper left", ncol=2)
+
+ ax.legend(legend, loc="upper left", ncol=2)
ax.set_xlabel("Frequency (GHz)")
ax.set_ylabel(label_list[leg] + " - " + component + " part")
ax.set_title(label_list[leg] + " - " + component + " part")
-
+
if zoom is True:
in_ax = inset_axes(ax,
- width="30%", # width = 30% of parent_bbox
- height=1.5, # height : 1 inch
- loc=1)
-
+ width="30%", # width = 30% of parent_bbox
+ height=1.5, # height : 1 inch
+ loc=1)
+
total_imp = 0
- for index in sorted_index:
+ for index in sorted_index:
attr = attr_list[index]
# Set all impedances with common indexes using + zero_impedance
try:
- sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
- in_ax.fill_between(sum_imp.data.index*1e-3, total_imp,
- total_imp + sum_imp.data[component]*1e-9, edgecolor=colorblind[index%10], color=colorblind[index%10])
+ sum_imp = getattr(getattr(self, attr),
+ Z_type) + zero_impedance
+ in_ax.fill_between(sum_imp.data.index*1e-3, total_imp,
+ total_imp + sum_imp.data[component]*1e-9, edgecolor=colorblind[index % 10], color=colorblind[index % 10])
total_imp += sum_imp.data[component]*1e-9
except AttributeError:
pass
in_ax.set_xlim([0, 200])
in_ax.set_xlabel("Frequency (kHz)")
in_ax.set_ylabel(r"$[\mathrm{G}\Omega]$")
-
+
return fig
-
- def effective_impedance(self, m, mu, sigma, M, tuneS, xi=None,
+
+ def effective_impedance(self, m, mu, sigma, M, tuneS, xi=None,
mode="Hermite"):
"""
Compute the longitudinal and transverse effective impedance.
@@ -359,45 +420,44 @@ class ImpedanceModel():
"""
attr_list = self.sum_names
-
- eff_array = np.zeros((len(attr_list),3), dtype=complex)
-
+
+ eff_array = np.zeros((len(attr_list), 3), dtype=complex)
+
for i, attr in enumerate(attr_list):
try:
impedance = getattr(getattr(self, attr), "Zlong")
- eff_array[i,0] = effective_impedance(self.ring, impedance,
- m, mu, sigma, M, tuneS,
- xi, mode)
+ eff_array[i, 0] = effective_impedance(self.ring, impedance,
+ m, mu, sigma, M, tuneS,
+ xi, mode)
except AttributeError:
pass
-
+
try:
impedance = getattr(getattr(self, attr), "Zxdip")
- eff_array[i,1] = effective_impedance(self.ring, impedance,
- m, mu, sigma, M, tuneS,
- xi, mode)
+ eff_array[i, 1] = effective_impedance(self.ring, impedance,
+ m, mu, sigma, M, tuneS,
+ xi, mode)
except AttributeError:
pass
-
+
try:
impedance = getattr(getattr(self, attr), "Zydip")
- eff_array[i,2] = effective_impedance(self.ring, impedance,
- m, mu, sigma, M, tuneS,
- xi, mode)
+ eff_array[i, 2] = effective_impedance(self.ring, impedance,
+ m, mu, sigma, M, tuneS,
+ xi, mode)
except AttributeError:
pass
-
- eff_array[:,0] = eff_array[:,0]*self.ring.omega0*1e3
- eff_array[:,1] = eff_array[:,1]*1e-3
- eff_array[:,2] = eff_array[:,2]*1e-3
-
- summary = pd.DataFrame(eff_array, index=attr_list,
- columns=["Z/n [mOhm]",
- "sum betax x Zxeff [kOhm]",
+
+ eff_array[:, 0] = eff_array[:, 0]*self.ring.omega0*1e3
+ eff_array[:, 1] = eff_array[:, 1]*1e-3
+ eff_array[:, 2] = eff_array[:, 2]*1e-3
+
+ summary = pd.DataFrame(eff_array, index=attr_list,
+ columns=["Z/n [mOhm]",
+ "sum betax x Zxeff [kOhm]",
"sum betay x Zyeff [kOhm]"])
-
- return summary
+ return summary
def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
"""
@@ -425,47 +485,51 @@ class ImpedanceModel():
full model and for each type of different component.
"""
-
+
fmax = self.sum.Zlong.data.index.max()
fmin = self.sum.Zlong.data.index.min()
-
+
Q = I*self.ring.T0/M
-
+
if fmin >= 0:
fmin = -1*fmax
f = np.linspace(fmin, fmax, int(n_points))
- beam_spect = beam_spectrum(f, M, bunch_spacing, sigma= sigma)
-
+ beam_spect = beam_spectrum(f, M, bunch_spacing, sigma=sigma)
+
bunch_spect = gaussian_bunch_spectrum(f, sigma)
-
+
attr_list = self.sum_names
-
- loss_array = np.zeros((len(attr_list),2))
-
+
+ loss_array = np.zeros((len(attr_list), 2))
+
for i, attr in enumerate(attr_list):
try:
impedance = getattr(getattr(self, attr), "Zlong")
- loss_array[i,0] = beam_loss_factor(impedance, f, beam_spect, self.ring)
- loss_array[i,1] = beam_loss_factor(impedance, f, bunch_spect, self.ring)
+ loss_array[i, 0] = beam_loss_factor(
+ impedance, f, beam_spect, self.ring)
+ loss_array[i, 1] = beam_loss_factor(
+ impedance, f, bunch_spect, self.ring)
except AttributeError:
pass
loss_array = loss_array*1e-12
- summary = pd.DataFrame(loss_array, index=attr_list,
+ summary = pd.DataFrame(loss_array, index=attr_list,
columns=["loss factor (beam) [V/pC]", "loss factor (bunch) [V/pC]"])
-
- summary["P (beam) [W]"] = summary["loss factor (beam) [V/pC]"]*1e12*Q**2/(self.ring.T0)
- summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"]*1e12*Q**2/(self.ring.T0)*M
-
+
+ summary["P (beam) [W]"] = summary["loss factor (beam) [V/pC]"] * \
+ 1e12*Q**2/(self.ring.T0)
+ summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"] * \
+ 1e12*Q**2/(self.ring.T0)*M
+
return summary
-
- def power_loss_spectrum(self, sigma, M, bunch_spacing, I, n_points=10e6,
- max_overlap=False,plot=False):
+
+ def power_loss_spectrum(self, sigma, M, bunch_spacing, I, n_points=10e6,
+ max_overlap=False, plot=False):
"""
Compute the power loss spectrum of the summed longitudinal impedance
as in Eq. (4) of [1].
-
+
Assume Gaussian bunches and constant spacing between bunches.
Parameters
@@ -498,7 +562,7 @@ class ImpedanceModel():
Frequency points.
power_loss : array
Power loss spectrum in [W].
-
+
References
----------
[1] : L. Teofili, et al. "A Multi-Physics Approach to Simulate the RF
@@ -506,60 +570,60 @@ class ImpedanceModel():
Device", in IPAC'18, 2018, doi:10.18429/JACoW-IPAC2018-THPAK093
"""
-
+
impedance = self.sum.Zlong
fmax = impedance.data.index.max()
fmin = impedance.data.index.min()
-
+
Q = I*self.ring.T0/M
-
+
if fmin >= 0:
fmin = -1*fmax
double_sided_impedance(impedance)
-
+
frequency = np.linspace(fmin, fmax, int(n_points))
if max_overlap is False:
- spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma)
+ spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma)
else:
spectrum = gaussian_bunch_spectrum(frequency, sigma)*M
-
+
pmax = np.floor(fmax/self.ring.f0)
pmin = np.floor(fmin/self.ring.f0)
-
- p = np.arange(pmin+1,pmax)
+
+ p = np.arange(pmin+1, pmax)
pf0 = p*self.ring.f0
ReZ = np.real(impedance(pf0))
spectral_density = np.abs(spectrum)**2
- # interpolation of the spectrum is needed to avoid problems liked to
+ # interpolation of the spectrum is needed to avoid problems liked to
# division by 0
# computing the spectrum directly to the frequency points gives
# wrong results
spect = interp1d(frequency, spectral_density)
power_loss = (self.ring.f0 * Q)**2 * ReZ * spect(pf0)
-
+
if plot is True:
fig, ax = plt.subplots()
twin = ax.twinx()
- p1, = ax.plot(pf0, ReZ, color="r",label="Re[Z] [Ohm]")
- p2, = twin.plot(pf0, spect(pf0)*(self.ring.f0 * Q)**2, color="b",
+ p1, = ax.plot(pf0, ReZ, color="r", label="Re[Z] [Ohm]")
+ p2, = twin.plot(pf0, spect(pf0)*(self.ring.f0 * Q)**2, color="b",
label="Beam spectrum [a.u.]")
ax.set_xlabel("Frequency [Hz]")
ax.set_ylabel("Re[Z] [Ohm]")
twin.set_ylabel("Beam spectrum [a.u.]")
plots = [p1, p2]
ax.legend(handles=plots, loc="best")
-
+
fig, ax = plt.subplots()
ax.plot(pf0, power_loss)
ax.set_xlabel("Frequency [Hz]")
ax.set_ylabel("Power loss [W]")
-
+
return pf0, power_loss
-
+
def save(self, file):
"""
Save impedance model to file.
-
+
The same pandas version is needed on both saving and loading computer
for the pickle to work.
@@ -573,18 +637,18 @@ class ImpedanceModel():
None.
"""
- to_save = {"wakefields":self.wakefields,
- "positions":self.positions,
- "names":self.names,
- "globals":self.globals,
- "globals_names":self.globals_names}
- with open(file,"wb") as f:
+ to_save = {"wakefields": self.wakefields,
+ "positions": self.positions,
+ "names": self.names,
+ "globals": self.globals,
+ "globals_names": self.globals_names}
+ with open(file, "wb") as f:
pickle.dump(to_save, f)
-
+
def load(self, file):
"""
Load impedance model from file.
-
+
The same pandas version is needed on both saving and loading computer
for the pickle to work.
@@ -598,10 +662,10 @@ class ImpedanceModel():
None.
"""
-
+
with open(file, 'rb') as f:
to_load = pickle.load(f)
-
+
self.wakefields = to_load["wakefields"]
self.positions = to_load["positions"]
self.names = to_load["names"]