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
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
mbtrack2/impedance/impedance_model.py
+ 195
− 131
View file @ 5ac810f4
...@@ -11,26 +11,28 @@ from copy import deepcopy ...@@ -11,26 +11,28 @@ from copy import deepcopy
from scipy.integrate import trapz from scipy.integrate import trapz
from scipy.interpolate import interp1d from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1.inset_locator import inset_axes 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) double_sided_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():
""" """
Define the impedance model of the machine. 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. 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). by the beta at the tracking location (ring.optics.local_beta).
Parameters Parameters
---------- ----------
ring : Synchrotron object ring : Synchrotron object
Attributes Attributes
---------- ----------
wakefields : list of WakeField objects wakefields : list of WakeField objects
...@@ -44,12 +46,12 @@ class ImpedanceModel(): ...@@ -44,12 +46,12 @@ class ImpedanceModel():
sum_"name" : WakeField sum_"name" : WakeField
Sum of the "name" Wakefield weigthed by beta functions. Sum of the "name" Wakefield weigthed by beta functions.
sum_names : array 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 : list of WakeField objects
Globals WakeFields in the model. Globals WakeFields in the model.
globals_names : list of str globals_names : list of str
Names of the global WakeFields objects. Names of the global WakeFields objects.
Methods Methods
------- -------
add(wakefield, positions, name) add(wakefield, positions, name)
...@@ -70,7 +72,7 @@ class ImpedanceModel(): ...@@ -70,7 +72,7 @@ class ImpedanceModel():
load(file) load(file)
Load impedance model from file. Load impedance model from file.
""" """
def __init__(self, ring): def __init__(self, ring):
self.ring = ring self.ring = ring
self.optics = self.ring.optics self.optics = self.ring.optics
...@@ -84,7 +86,7 @@ class ImpedanceModel(): ...@@ -84,7 +86,7 @@ class ImpedanceModel():
def add(self, wakefield, positions, name=None): def add(self, wakefield, positions, name=None):
""" """
Add the same WakeField object at different locations to the model. Add the same WakeField object at different locations to the model.
Parameters Parameters
---------- ----------
wakefield : WakeField wakefield : WakeField
...@@ -92,7 +94,7 @@ class ImpedanceModel(): ...@@ -92,7 +94,7 @@ class ImpedanceModel():
positions : array, float or int positions : array, float or int
Array of longitudinal positions where the elements are loacted. Array of longitudinal positions where the elements are loacted.
name : str, optional 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. is used. The default is None.
Returns Returns
...@@ -110,22 +112,22 @@ class ImpedanceModel(): ...@@ -110,22 +112,22 @@ class ImpedanceModel():
self.names.append(name) self.names.append(name)
else: else:
raise ValueError("This name is already taken.") raise ValueError("This name is already taken.")
def add_global(self, wakefield, name=None): def add_global(self, wakefield, name=None):
""" """
Add a "global" WakeField object which will added to the sum WakeField Add a "global" WakeField object which will added to the sum WakeField
but not weighted by beta functions. but not weighted by beta functions.
To use with "distributed" elements, for example a resistive wall To use with "distributed" elements, for example a resistive wall
wakefield computed from an effective radius (so which has already been wakefield computed from an effective radius (so which has already been
weighted by beta functions). weighted by beta functions).
Parameters Parameters
---------- ----------
wakefield : WakeField wakefield : WakeField
WakeField object to add to the model. WakeField object to add to the model.
name : str, optional 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. is used. The default is None.
Returns Returns
...@@ -143,7 +145,7 @@ class ImpedanceModel(): ...@@ -143,7 +145,7 @@ class ImpedanceModel():
else: else:
raise ValueError("This name is already taken.") raise ValueError("This name is already taken.")
# setattr(self, name, wakefield) # setattr(self, name, wakefield)
def sum_beta(self, wake, beta): def sum_beta(self, wake, beta):
""" """
Weight a WakeField object by an array of beta functions. Weight a WakeField object by an array of beta functions.
...@@ -165,8 +167,8 @@ class ImpedanceModel(): ...@@ -165,8 +167,8 @@ class ImpedanceModel():
local_beta = self.ring.optics.local_beta local_beta = self.ring.optics.local_beta
for component_name in wake_sum.components: for component_name in wake_sum.components:
comp = getattr(wake_sum, component_name) comp = getattr(wake_sum, component_name)
weight = ((beta[0,:] ** comp.power_x) * weight = ((beta[0, :] ** comp.power_x) *
(beta[1,:] ** comp.power_y)) (beta[1, :] ** comp.power_y))
if comp.plane == "x": if comp.plane == "x":
weight = weight.sum() / local_beta[0] weight = weight.sum() / local_beta[0]
if comp.plane == "y": if comp.plane == "y":
...@@ -175,7 +177,7 @@ class ImpedanceModel(): ...@@ -175,7 +177,7 @@ class ImpedanceModel():
weight = weight.sum() weight = weight.sum()
setattr(wake_sum, component_name, weight*comp) setattr(wake_sum, component_name, weight*comp)
return wake_sum return wake_sum
def compute_sum_names(self): def compute_sum_names(self):
""" """
Compute the weighted WakeField for each WakeField object type. Compute the weighted WakeField for each WakeField object type.
...@@ -188,19 +190,19 @@ class ImpedanceModel(): ...@@ -188,19 +190,19 @@ class ImpedanceModel():
wake_sum.name = attribute_name wake_sum.name = attribute_name
setattr(self, attribute_name, wake_sum) setattr(self, attribute_name, wake_sum)
self.sum_names.append(attribute_name) self.sum_names.append(attribute_name)
def compute_sum(self): def compute_sum(self):
""" """
Compute the sum of all weighted WakeField into self.sum. Compute the sum of all weighted WakeField into self.sum.
""" """
self.compute_sum_names() self.compute_sum_names()
for i, name in enumerate(self.sum_names): for i, name in enumerate(self.sum_names):
if i==0: if i == 0:
self.sum = deepcopy(getattr(self, name)) self.sum = deepcopy(getattr(self, name))
self.sum.name = "sum" self.sum.name = "sum"
else: else:
wake2 = getattr(self, name) wake2 = getattr(self, name)
for component_name2 in wake2.components: for component_name2 in wake2.components:
comp2 = getattr(wake2, component_name2) comp2 = getattr(wake2, component_name2)
try: try:
comp1 = getattr(self.sum, component_name2) comp1 = getattr(self.sum, component_name2)
...@@ -215,7 +217,7 @@ class ImpedanceModel(): ...@@ -215,7 +217,7 @@ class ImpedanceModel():
self.sum = deepcopy(wake2) self.sum = deepcopy(wake2)
self.sum.name = "sum" self.sum.name = "sum"
else: else:
for component_name2 in wake2.components: for component_name2 in wake2.components:
comp2 = getattr(wake2, component_name2) comp2 = getattr(wake2, component_name2)
try: try:
comp1 = getattr(self.sum, component_name2) comp1 = getattr(self.sum, component_name2)
...@@ -223,7 +225,64 @@ class ImpedanceModel(): ...@@ -223,7 +225,64 @@ class ImpedanceModel():
except AttributeError: except AttributeError:
setattr(self.sum, component_name2, comp2) 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): attr_list=None, zoom=False):
""" """
Plot the contributions of different kind of WakeFields. Plot the contributions of different kind of WakeFields.
...@@ -246,46 +305,47 @@ class ImpedanceModel(): ...@@ -246,46 +305,47 @@ 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})$",
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_{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_{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})$"] 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] leg = Ztype_dict[Z_type]
# sort plot by decresing area size # sort plot by decresing area size
area = np.zeros((len(attr_list),)) area = np.zeros((len(attr_list),))
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]
# Init fig # Init fig
fig = plt.figure() fig = plt.figure()
ax = fig.gca() ax = fig.gca()
zero_impedance = getattr(self.sum, Z_type)*0 zero_impedance = getattr(self.sum, Z_type)*0
total_imp = 0 total_imp = 0
legend = [] legend = []
if sigma is not None: if sigma is not None:
legend.append("Spectral density for sigma = " + str(sigma) + " s") legend.append("Spectral density for sigma = " + str(sigma) + " s")
# Main plot # Main plot
colorblind = colormaps['tab10'].colors colorblind = colormaps['tab10'].colors
for index in sorted_index: for index in sorted_index:
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
ax.fill_between(sum_imp.data.index*1e-9, total_imp, 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], edgecolor=colorblind[index % 10], color=colorblind[index % 10])
total_imp += sum_imp.data[component]*scale[leg] total_imp += sum_imp.data[component]*scale[leg]
if attr[:4] == "sum_": if attr[:4] == "sum_":
legend.append(attr[4:]) legend.append(attr[4:])
...@@ -293,41 +353,42 @@ class ImpedanceModel(): ...@@ -293,41 +353,42 @@ class ImpedanceModel():
legend.append(attr) legend.append(attr)
except AttributeError: except AttributeError:
pass pass
if sigma is not None: if sigma is not None:
spect = spectral_density(zero_impedance.data.index, sigma) spect = spectral_density(zero_impedance.data.index, sigma)
spect = spect/spect.max()*total_imp.max() spect = spect/spect.max()*total_imp.max()
ax.plot(sum_imp.data.index*1e-9, spect, 'r', linewidth=2.5) 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_xlabel("Frequency (GHz)")
ax.set_ylabel(label_list[leg] + " - " + component + " part") ax.set_ylabel(label_list[leg] + " - " + component + " part")
ax.set_title(label_list[leg] + " - " + component + " part") ax.set_title(label_list[leg] + " - " + component + " part")
if zoom is True: if zoom is True:
in_ax = inset_axes(ax, in_ax = inset_axes(ax,
width="30%", # width = 30% of parent_bbox width="30%", # width = 30% of parent_bbox
height=1.5, # height : 1 inch height=1.5, # height : 1 inch
loc=1) loc=1)
total_imp = 0 total_imp = 0
for index in sorted_index: for index in sorted_index:
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),
in_ax.fill_between(sum_imp.data.index*1e-3, total_imp, Z_type) + zero_impedance
total_imp + sum_imp.data[component]*1e-9, edgecolor=colorblind[index%10], color=colorblind[index%10]) 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 total_imp += sum_imp.data[component]*1e-9
except AttributeError: except AttributeError:
pass pass
in_ax.set_xlim([0, 200]) in_ax.set_xlim([0, 200])
in_ax.set_xlabel("Frequency (kHz)") in_ax.set_xlabel("Frequency (kHz)")
in_ax.set_ylabel(r"$[\mathrm{G}\Omega]$") in_ax.set_ylabel(r"$[\mathrm{G}\Omega]$")
return fig 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"): mode="Hermite"):
""" """
Compute the longitudinal and transverse effective impedance. Compute the longitudinal and transverse effective impedance.
...@@ -359,45 +420,44 @@ class ImpedanceModel(): ...@@ -359,45 +420,44 @@ class ImpedanceModel():
""" """
attr_list = self.sum_names 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): for i, attr in enumerate(attr_list):
try: try:
impedance = getattr(getattr(self, attr), "Zlong") impedance = getattr(getattr(self, attr), "Zlong")
eff_array[i,0] = effective_impedance(self.ring, impedance, eff_array[i, 0] = effective_impedance(self.ring, impedance,
m, mu, sigma, M, tuneS, m, mu, sigma, M, tuneS,
xi, mode) xi, mode)
except AttributeError: except AttributeError:
pass pass
try: try:
impedance = getattr(getattr(self, attr), "Zxdip") impedance = getattr(getattr(self, attr), "Zxdip")
eff_array[i,1] = effective_impedance(self.ring, impedance, eff_array[i, 1] = effective_impedance(self.ring, impedance,
m, mu, sigma, M, tuneS, m, mu, sigma, M, tuneS,
xi, mode) xi, mode)
except AttributeError: except AttributeError:
pass pass
try: try:
impedance = getattr(getattr(self, attr), "Zydip") impedance = getattr(getattr(self, attr), "Zydip")
eff_array[i,2] = effective_impedance(self.ring, impedance, eff_array[i, 2] = effective_impedance(self.ring, impedance,
m, mu, sigma, M, tuneS, m, mu, sigma, M, tuneS,
xi, mode) xi, mode)
except AttributeError: except AttributeError:
pass pass
eff_array[:,0] = eff_array[:,0]*self.ring.omega0*1e3 eff_array[:, 0] = eff_array[:, 0]*self.ring.omega0*1e3
eff_array[:,1] = eff_array[:,1]*1e-3 eff_array[:, 1] = eff_array[:, 1]*1e-3
eff_array[:,2] = eff_array[:,2]*1e-3 eff_array[:, 2] = eff_array[:, 2]*1e-3
summary = pd.DataFrame(eff_array, index=attr_list, summary = pd.DataFrame(eff_array, index=attr_list,
columns=["Z/n [mOhm]", columns=["Z/n [mOhm]",
"sum betax x Zxeff [kOhm]", "sum betax x Zxeff [kOhm]",
"sum betay x Zyeff [kOhm]"]) "sum betay x Zyeff [kOhm]"])
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):
""" """
...@@ -425,47 +485,51 @@ class ImpedanceModel(): ...@@ -425,47 +485,51 @@ class ImpedanceModel():
full model and for each type of different component. full model and for each type of different component.
""" """
fmax = self.sum.Zlong.data.index.max() fmax = self.sum.Zlong.data.index.max()
fmin = self.sum.Zlong.data.index.min() fmin = self.sum.Zlong.data.index.min()
Q = I*self.ring.T0/M Q = I*self.ring.T0/M
if fmin >= 0: if fmin >= 0:
fmin = -1*fmax fmin = -1*fmax
f = np.linspace(fmin, fmax, int(n_points)) 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) bunch_spect = gaussian_bunch_spectrum(f, sigma)
attr_list = self.sum_names 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): 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
loss_array = loss_array*1e-12 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]"]) 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
def power_loss_spectrum(self, sigma, M, bunch_spacing, I, n_points=10e6, def power_loss_spectrum(self, sigma, M, bunch_spacing, I, n_points=10e6,
max_overlap=False,plot=False): max_overlap=False, plot=False):
""" """
Compute the power loss spectrum of the summed longitudinal impedance Compute the power loss spectrum of the summed longitudinal impedance
as in Eq. (4) of [1]. as in Eq. (4) of [1].
Assume Gaussian bunches and constant spacing between bunches. Assume Gaussian bunches and constant spacing between bunches.
Parameters Parameters
...@@ -498,7 +562,7 @@ class ImpedanceModel(): ...@@ -498,7 +562,7 @@ class ImpedanceModel():
Frequency points. Frequency points.
power_loss : array power_loss : array
Power loss spectrum in [W]. Power loss spectrum in [W].
References References
---------- ----------
[1] : L. Teofili, et al. "A Multi-Physics Approach to Simulate the RF [1] : L. Teofili, et al. "A Multi-Physics Approach to Simulate the RF
...@@ -506,60 +570,60 @@ class ImpedanceModel(): ...@@ -506,60 +570,60 @@ class ImpedanceModel():
Device", in IPAC'18, 2018, doi:10.18429/JACoW-IPAC2018-THPAK093 Device", in IPAC'18, 2018, doi:10.18429/JACoW-IPAC2018-THPAK093
""" """
impedance = self.sum.Zlong impedance = self.sum.Zlong
fmax = impedance.data.index.max() fmax = impedance.data.index.max()
fmin = impedance.data.index.min() fmin = impedance.data.index.min()
Q = I*self.ring.T0/M Q = I*self.ring.T0/M
if fmin >= 0: if fmin >= 0:
fmin = -1*fmax fmin = -1*fmax
double_sided_impedance(impedance) double_sided_impedance(impedance)
frequency = np.linspace(fmin, fmax, int(n_points)) frequency = np.linspace(fmin, fmax, int(n_points))
if max_overlap is False: if max_overlap is False:
spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma) spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma)
else: else:
spectrum = gaussian_bunch_spectrum(frequency, sigma)*M spectrum = gaussian_bunch_spectrum(frequency, sigma)*M
pmax = np.floor(fmax/self.ring.f0) pmax = np.floor(fmax/self.ring.f0)
pmin = np.floor(fmin/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 pf0 = p*self.ring.f0
ReZ = np.real(impedance(pf0)) ReZ = np.real(impedance(pf0))
spectral_density = np.abs(spectrum)**2 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 # division by 0
# computing the spectrum directly to the frequency points gives # computing the spectrum directly to the frequency points gives
# wrong results # wrong results
spect = interp1d(frequency, spectral_density) spect = interp1d(frequency, spectral_density)
power_loss = (self.ring.f0 * Q)**2 * ReZ * spect(pf0) power_loss = (self.ring.f0 * Q)**2 * ReZ * spect(pf0)
if plot is True: if plot is True:
fig, ax = plt.subplots() fig, ax = plt.subplots()
twin = ax.twinx() twin = ax.twinx()
p1, = ax.plot(pf0, ReZ, color="r",label="Re[Z] [Ohm]") p1, = ax.plot(pf0, ReZ, color="r", label="Re[Z] [Ohm]")
p2, = twin.plot(pf0, spect(pf0)*(self.ring.f0 * Q)**2, color="b", p2, = twin.plot(pf0, spect(pf0)*(self.ring.f0 * Q)**2, color="b",
label="Beam spectrum [a.u.]") label="Beam spectrum [a.u.]")
ax.set_xlabel("Frequency [Hz]") ax.set_xlabel("Frequency [Hz]")
ax.set_ylabel("Re[Z] [Ohm]") ax.set_ylabel("Re[Z] [Ohm]")
twin.set_ylabel("Beam spectrum [a.u.]") twin.set_ylabel("Beam spectrum [a.u.]")
plots = [p1, p2] plots = [p1, p2]
ax.legend(handles=plots, loc="best") ax.legend(handles=plots, loc="best")
fig, ax = plt.subplots() fig, ax = plt.subplots()
ax.plot(pf0, power_loss) ax.plot(pf0, power_loss)
ax.set_xlabel("Frequency [Hz]") ax.set_xlabel("Frequency [Hz]")
ax.set_ylabel("Power loss [W]") ax.set_ylabel("Power loss [W]")
return pf0, power_loss return pf0, power_loss
def save(self, file): def save(self, file):
""" """
Save impedance model to file. Save impedance model to file.
The same pandas version is needed on both saving and loading computer The same pandas version is needed on both saving and loading computer
for the pickle to work. for the pickle to work.
...@@ -573,18 +637,18 @@ class ImpedanceModel(): ...@@ -573,18 +637,18 @@ class ImpedanceModel():
None. None.
""" """
to_save = {"wakefields":self.wakefields, to_save = {"wakefields": self.wakefields,
"positions":self.positions, "positions": self.positions,
"names":self.names, "names": self.names,
"globals":self.globals, "globals": self.globals,
"globals_names":self.globals_names} "globals_names": self.globals_names}
with open(file,"wb") as f: with open(file, "wb") as f:
pickle.dump(to_save, f) pickle.dump(to_save, f)
def load(self, file): def load(self, file):
""" """
Load impedance model from file. Load impedance model from file.
The same pandas version is needed on both saving and loading computer The same pandas version is needed on both saving and loading computer
for the pickle to work. for the pickle to work.
...@@ -598,10 +662,10 @@ class ImpedanceModel(): ...@@ -598,10 +662,10 @@ class ImpedanceModel():
None. None.
""" """
with open(file, 'rb') as f: with open(file, 'rb') as f:
to_load = pickle.load(f) to_load = pickle.load(f)
self.wakefields = to_load["wakefields"] self.wakefields = to_load["wakefields"]
self.positions = to_load["positions"] self.positions = to_load["positions"]
self.names = to_load["names"] self.names = to_load["names"]
......
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