diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index 2985cb784e8c79bca6ea6718ff3bfe7730af5636..e23b17ad0caa380f1a049ab6a7059aa7a3fd1e8a 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -6,6 +6,7 @@ import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pickle
+from copy import deepcopy
from scipy.integrate import trapz
from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
@@ -14,52 +15,43 @@ from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance,
from mbtrack2.utilities.spectrum import (beam_spectrum,
gaussian_bunch_spectrum,
spectral_density)
-from mbtrack2.impedance.wakefield import WakeField
-from mbtrack2.tracking.element import Element
-class ImpedanceModel(Element):
+class ImpedanceModel():
"""
Define the impedance model of the machine.
+ The model must be completed with successive add(...) calls, then
+ compute_sum() must be run.
+
Parameters
----------
ring : Synchrotron object
- wakefield_list : list of WakeField objects
- WakeFields to add to the model.
- wakefiled_positions : list
- Longitudinal positions corresponding to the added Wakfields.
Attributes
----------
wakefields : list of WakeField objects
WakeFields in the model.
- positions : array
- WakeFields positions.
- names : array
- Names of (unique) WakeField objects.
+ positions : list of arrays
+ Positions corresponding the different WakeFields.
+ names : list of str
+ Names of the WakeField objects.
sum : WakeField
- Sum of every WakeField in the model.
+ Sum of every WakeField in the model weigthed by beta functions.
sum_"name" : WakeField
- Sum of every Wakefield with the same "name".
+ Sum of the "name" Wakefield weigthed by beta functions.
sum_names : array
Names of attributes where the WakeFields are summed by name.
Methods
-------
- add(wakefield_list, wakefiled_positions)
- Add elements to the model.
- add_multiple_elements(wakefield, wakefiled_positions)
- Add the same element at different locations to the model.
- sum_elements()
- Sum all WakeFields into self.sum.
- update_name_list()
- Update self.names with uniques names of self.wakefields.
- find_wakefield(name)
- Return indexes of WakeFields with the same name in self.wakefields.
- sum_by_name(name)
- Sum the elements with the same name in the model into sum_name.
- sum_by_name_all(name)
- Sum all the elements with the same name in the model into sum_name.
+ add(wakefield, positions, name)
+ Add the same WakeField object at different locations to the model.
+ sum_beta(wake, beta)
+ Weight a WakeField object by an array of beta functions.
+ compute_sum_names()
+ Compute the weighted WakeField for each WakeField object type.
+ compute_sum()
+ Compute the sum of all weighted WakeField into self.sum.
plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
Plot the contributions of different kind of WakeFields.
save(file)
@@ -68,132 +60,101 @@ class ImpedanceModel(Element):
Load impedance model from file.
"""
- def __init__(self, ring, wakefield_list=None, wakefiled_positions=None):
+ def __init__(self, ring):
self.ring = ring
self.optics = self.ring.optics
self.wakefields = []
- self.positions = np.array([])
- self.names = np.array([])
- self.sum_names = np.array([])
- self.add(wakefield_list, wakefiled_positions)
-
- def track(self, beam):
+ self.positions = []
+ self.names = []
+ self.sum_names = []
+
+ def add(self, wakefield, positions, name=None):
"""
- Track a beam object through this Element.
+ Add the same WakeField object at different locations to the model.
Parameters
----------
- beam : Beam object
+ wakefield : WakeField
+ WakeField object to add to the model.
+ 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
+ is used. The default is None.
+
+ Returns
+ -------
+ None.
+
"""
- raise NotImplementedError
-
- def sum_elements(self):
- """Sum all WakeFields into self.sum"""
- beta = self.optics.beta(self.positions)
- self.sum = WakeField.add_several_wakefields(self.wakefields, beta)
- if "sum" not in self.sum_names:
- self.sum_names = np.append(self.sum_names, "sum")
-
- def update_name_list(self):
- """Update self.names with uniques names of self.wakefields."""
- for wakefield in self.wakefields:
- if wakefield.name is None:
- continue
- if wakefield.name not in self.names:
- self.names = np.append(self.names, wakefield.name)
-
- def count_elements(self):
- """Count number of each type of WakeField in the model."""
- self.count = np.zeros(len(self.names))
- for wakefield in self.wakefields:
- if wakefield.name is not None:
- self.count += (wakefield.name == self.names).astype(int)
-
- def find_wakefield(self, name):
+ self.wakefields.append(wakefield)
+ self.positions.append(positions)
+ if name is None:
+ name = wakefield.name
+ if name is None:
+ raise ValueError("Please give a valid name.")
+ if name not in self.names:
+ self.names.append(name)
+ else:
+ raise ValueError("This name is already taken.")
+
+ @staticmethod
+ def sum_beta(wake, beta):
"""
- Return indexes of WakeFields with the same name in self.wakefields.
+ Weight a WakeField object by an array of beta functions.
Parameters
----------
- name : str
- WakeField name.
+ wake : WakeField
+ WakeField element object.
+ beta : array of shape (2, N)
+ Beta function at the locations of the elements.
Returns
-------
- index : list
- Index of positions in self.wakefields.
+ wake_sum : WakeField
+ WakeField object weighted by beta functions.
"""
- index = []
- for i, wakefield in enumerate(self.wakefields):
- if wakefield.name == name:
- index.append(i)
- return index
-
- def sum_by_name(self, name):
- """
- Sum the elements with the same name in the model into sum_name.
-
- Parameters
- ----------
- name : str
- Name of the WakeField to sum.
- """
- attribute_name = "sum_" + name
- index = self.find_wakefield(name)
- beta = self.optics.beta(self.positions[index])
- wakes = []
- for i in index:
- wakes.append(self.wakefields[i])
- wake_sum = WakeField.add_several_wakefields(wakes, beta)
- setattr(self, attribute_name, wake_sum)
- if attribute_name not in self.sum_names:
- self.sum_names = np.append(self.sum_names, attribute_name)
-
- def sum_by_name_all(self):
- """
- Sum all the elements with the same name in the model into sum_name.
+ wake_sum = deepcopy(wake)
+ for component_name in wake.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)
+ return wake_sum
+
+ def compute_sum_names(self):
"""
- for name in self.names:
- self.sum_by_name(name)
-
- def add(self, wakefield_list, wakefiled_positions):
+ Compute the weighted WakeField for each WakeField object type.
+ The new summed WakeField object are set to into self.sum_name.
"""
- Add elements to the model.
-
- Parameters
- ----------
- wakefield_list : list of WakeField objects
- WakeFields to add to the model.
- wakefiled_positions : list
- Longitudinal positions corresponding to the added Wakfields.
+ for idx, wake in enumerate(self.wakefields):
+ attribute_name = "sum_" + self.names[idx]
+ beta = self.optics.beta(self.positions[idx])
+ wake_sum = self.sum_beta(wake, beta)
+ setattr(self, attribute_name, wake_sum)
+ self.sum_names.append(attribute_name)
+
+ def compute_sum(self):
"""
- if (wakefield_list is not None) and (wakefiled_positions is not None):
- for wakefield in wakefield_list:
- self.wakefields.append(wakefield)
-
- for position in wakefiled_positions:
- self.positions = np.append(self.positions, position)
-
- self.update_name_list()
-
- def add_multiple_elements(self, wakefield, wakefiled_positions):
+ Compute the sum of all weighted WakeField into self.sum.
+ self.compute_sum_names must be called before this.
"""
- Add the same element at different locations to the model.
+ self.compute_sum_names()
+ for i, name in enumerate(self.sum_names):
+ if i==0:
+ self.sum = deepcopy(getattr(self, name))
+ else:
+ wake2 = getattr(self, name)
+ 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)
- Parameters
- ----------
- WakeField : WakeField object
- WakeField to add to the model.
- wakefiled_positions : list
- Longitudinal positions corresponding to the added Wakfield.
- """
- for position in wakefiled_positions:
- self.positions = np.append(self.positions, position)
- self.wakefields.append(wakefield)
-
- self.update_name_list()
-
def plot_area(self, Z_type="Zlong", component="real", sigma=None,
attr_list=None, zoom=False):
"""