From 1bf60a32cb8e4319273fc8e53bf1700f6e8003e5 Mon Sep 17 00:00:00 2001 From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr> Date: Fri, 21 Apr 2023 22:59:09 +0200 Subject: [PATCH] Rework ImpedanceModel The ImpedanceModel class can now deal with "global impedance" such a resistive wall in addition to do the beta weigthing for "localized" impedances. The summed wake function after beta weigthing is now normalized by the tracking beta function (ring.optics.local_beta). --- mbtrack2/impedance/impedance_model.py | 102 ++++++++++++++++++++++---- 1 file changed, 87 insertions(+), 15 deletions(-) diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py index f97d625..223b634 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() -- GitLab