Skip to content
Snippets Groups Projects
Commit 5ac810f4 authored by Vadim Gubaidulin's avatar Vadim Gubaidulin
Browse files

Added two utility functions:

1) group_attributes() will group attributes in the impedance model.
A group can be formed from a list of elements.
This list can be autogenerated by a matching string
2) rename_attribute() simple wrapper to rename an attribute.
Can be used to rename impedance model components.
parent d4b1c2cc
Branches
Tags
No related merge requests found
...@@ -16,6 +16,8 @@ from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance, ...@@ -16,6 +16,8 @@ from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance,
from mbtrack2.utilities.spectrum import (beam_spectrum, from mbtrack2.utilities.spectrum import (beam_spectrum,
gaussian_bunch_spectrum, gaussian_bunch_spectrum,
spectral_density) spectral_density)
from mbtrack2.impedance.wakefield import WakeField, WakeFunction
class ImpedanceModel(): class ImpedanceModel():
""" """
...@@ -223,6 +225,63 @@ class ImpedanceModel(): ...@@ -223,6 +225,63 @@ class ImpedanceModel():
except AttributeError: except AttributeError:
setattr(self.sum, component_name2, comp2) setattr(self.sum, component_name2, comp2)
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, def plot_area(self, Z_type="Zlong", component="real", sigma=None,
attr_list=None, zoom=False): attr_list=None, zoom=False):
""" """
...@@ -246,9 +305,9 @@ class ImpedanceModel(): ...@@ -246,9 +305,9 @@ class ImpedanceModel():
if attr_list is None: if attr_list is None:
attr_list = self.sum_names attr_list = self.sum_names
# manage legend # 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] scale = [1e-3, 1e-6, 1e-6, 1e-6, 1e-6]
label_list = [r"$Z_\mathrm{long} \; (\mathrm{k}\Omega)$", 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_{x,j} Z_{x,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
...@@ -262,7 +321,8 @@ class ImpedanceModel(): ...@@ -262,7 +321,8 @@ class ImpedanceModel():
for index, attr in enumerate(attr_list): for index, attr in enumerate(attr_list):
try: try:
sum_imp = getattr(getattr(self, attr), Z_type) 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: except AttributeError:
pass pass
sorted_index = area.argsort()[::-1] sorted_index = area.argsort()[::-1]
...@@ -315,7 +375,8 @@ class ImpedanceModel(): ...@@ -315,7 +375,8 @@ class ImpedanceModel():
attr = attr_list[index] attr = attr_list[index]
# Set all impedances with common indexes using + zero_impedance # Set all impedances with common indexes using + zero_impedance
try: try:
sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance sum_imp = getattr(getattr(self, attr),
Z_type) + zero_impedance
in_ax.fill_between(sum_imp.data.index*1e-3, total_imp, 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, edgecolor=colorblind[index % 10], color=colorblind[index % 10])
total_imp += sum_imp.data[component]*1e-9 total_imp += sum_imp.data[component]*1e-9
...@@ -398,7 +459,6 @@ class ImpedanceModel(): ...@@ -398,7 +459,6 @@ class ImpedanceModel():
return summary return summary
def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6): def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
""" """
Compute the beam and bunch loss factor and energy losses for each type Compute the beam and bunch loss factor and energy losses for each type
...@@ -446,8 +506,10 @@ class ImpedanceModel(): ...@@ -446,8 +506,10 @@ class ImpedanceModel():
for i, attr in enumerate(attr_list): for i, attr in enumerate(attr_list):
try: try:
impedance = getattr(getattr(self, attr), "Zlong") impedance = getattr(getattr(self, attr), "Zlong")
loss_array[i,0] = beam_loss_factor(impedance, f, beam_spect, self.ring) loss_array[i, 0] = beam_loss_factor(
loss_array[i,1] = beam_loss_factor(impedance, f, bunch_spect, self.ring) impedance, f, beam_spect, self.ring)
loss_array[i, 1] = beam_loss_factor(
impedance, f, bunch_spect, self.ring)
except AttributeError: except AttributeError:
pass pass
...@@ -455,8 +517,10 @@ class ImpedanceModel(): ...@@ -455,8 +517,10 @@ class ImpedanceModel():
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]"]) 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 (beam) [W]"] = summary["loss factor (beam) [V/pC]"] * \
summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"]*1e12*Q**2/(self.ring.T0)*M 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 return summary
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment