diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index f97d6256ab62f16fec7cf60944f25bc4f376fd09..223b634f05d0812858cdaeb6f8f1421cfc30aef6 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -20,8 +20,11 @@ class ImpedanceModel():
"""
Define the impedance model of the machine.
- The model must be completed with successive add(...) calls, then
- compute_sum() must be run.
+ 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
+ by the beta at the tracking location (ring.optics.local_beta).
Parameters
----------
@@ -40,12 +43,19 @@ 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)
Add the same WakeField object at different locations to the model.
+ add_global(wakefield, name)
+ Add a "global" WakeField object which will added to the sum WakeField
+ but not weighted by beta functions.
sum_beta(wake, beta)
Weight a WakeField object by an array of beta functions.
compute_sum_names()
@@ -66,6 +76,8 @@ class ImpedanceModel():
self.wakefields = []
self.positions = []
self.names = []
+ self.globals = []
+ self.globals_names = []
self.sum_names = []
def add(self, wakefield, positions, name=None):
@@ -97,9 +109,41 @@ 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.
- @staticmethod
- def sum_beta(wake, beta):
+ 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
+ is used. The default is None.
+
+ Returns
+ -------
+ None.
+
+ """
+ self.globals.append(wakefield)
+ if name is None:
+ name = wakefield.name
+ if name is None:
+ raise ValueError("Please give a valid name.")
+ if name not in self.globals_names:
+ self.globals_names.append(name)
+ 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.
@@ -117,11 +161,18 @@ class ImpedanceModel():
"""
wake_sum = deepcopy(wake)
- for component_name in wake.components:
+ 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))
- setattr(wake_sum, component_name, weight.sum()*comp)
+ if comp.plane == "x":
+ weight = weight.sum() / local_beta[0]
+ if comp.plane == "y":
+ weight = weight.sum() / local_beta[1]
+ else:
+ weight = weight.sum()
+ setattr(wake_sum, component_name, weight*comp)
return wake_sum
def compute_sum_names(self):
@@ -133,18 +184,19 @@ class ImpedanceModel():
attribute_name = "sum_" + self.names[idx]
beta = self.optics.beta(self.positions[idx])
wake_sum = self.sum_beta(wake, beta)
+ 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 must be called before this.
"""
self.compute_sum_names()
for i, name in enumerate(self.sum_names):
if i==0:
self.sum = deepcopy(getattr(self, name))
+ self.sum.name = "sum"
else:
wake2 = getattr(self, name)
for component_name2 in wake2.components:
@@ -154,6 +206,21 @@ class ImpedanceModel():
setattr(self.sum, component_name2, comp1 + comp2)
except AttributeError:
setattr(self.sum, component_name2, comp2)
+ for i, wake2 in enumerate(self.globals):
+ name = self.globals_names[i]
+ setattr(self, name, wake2)
+ self.sum_names.append(name)
+ if not hasattr(self, "sum"):
+ self.sum = deepcopy(wake2)
+ self.sum.name = "sum"
+ else:
+ for component_name2 in wake2.components:
+ comp2 = getattr(wake2, component_name2)
+ try:
+ comp1 = getattr(self.sum, component_name2)
+ setattr(self.sum, component_name2, comp1 + comp2)
+ except AttributeError:
+ setattr(self.sum, component_name2, comp2)
def plot_area(self, Z_type="Zlong", component="real", sigma=None,
attr_list=None, zoom=False):
@@ -176,16 +243,17 @@ class ImpedanceModel():
"""
if attr_list is None:
- attr_list = self.sum_names[self.sum_names != "sum"]
+ attr_list = self.sum_names
+
# manage legend
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_{long} \; [k\Omega]$",
- r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega]$",
- r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega]$",
- r"$\sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega]$",
- r"$\sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega]$"]
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^{Dip} \; [M\Omega/m]$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^{Dip} \; [M\Omega/m]$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^{Quad} \; [M\Omega/m]$",
+ r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^{Quad} \; [M\Omega/m]$"]
leg = Ztype_dict[Z_type]
# sort plot by decresing area size
@@ -229,7 +297,7 @@ class ImpedanceModel():
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")
+ ax.legend(legend, loc="upper left", ncol=3)
ax.set_xlabel("Frequency [GHz]")
ax.set_ylabel(label_list[leg] + " - " + component + " part")
ax.set_title(label_list[leg] + " - " + component + " part")
@@ -505,7 +573,9 @@ class ImpedanceModel():
"""
to_save = {"wakefields":self.wakefields,
"positions":self.positions,
- "names":self.names}
+ "names":self.names,
+ "globals":self.globals,
+ "globals_names":self.globals_names}
with open(file,"wb") as f:
pickle.dump(to_save, f)
@@ -533,4 +603,6 @@ class ImpedanceModel():
self.wakefields = to_load["wakefields"]
self.positions = to_load["positions"]
self.names = to_load["names"]
+ self.globals = to_load["globals"]
+ self.globals_names = to_load["globals_names"]
self.compute_sum()