diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index 2985cb784e8c79bca6ea6718ff3bfe7730af5636..4fc135240c20b19c36c7189325712bc8cdb1a720 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):
"""
@@ -546,7 +507,8 @@ class ImpedanceModel(Element):
"""
to_save = {"wakefields":self.wakefields,
- "positions":self.positions}
+ "positions":self.positions,
+ "names":self.names}
with open(file,"wb") as f:
pickle.dump(to_save, f)
@@ -572,6 +534,6 @@ class ImpedanceModel(Element):
to_load = pickle.load(f)
self.wakefields = to_load["wakefields"]
- self.positions = to_load["positions"]
- self.sum_elements()
- self.sum_by_name_all()
+ self.positions = to_load["positions"]
+ self.names = to_load["names"]
+ self.compute_sum()
diff --git a/mbtrack2/utilities/optics.py b/mbtrack2/utilities/optics.py
index 6552918fb8eebc751c1f48ab4476d190ba7b8221..bcb3ed0d6eef319c743f654ae51febb8773bbada 100644
--- a/mbtrack2/utilities/optics.py
+++ b/mbtrack2/utilities/optics.py
@@ -554,8 +554,9 @@ class PhysicalModel:
return (1/sigma_star, a1_L, a2_L, a3_H, a4_H, a3_V, a4_V)
- def change_values(self, start_position, end_position, x_right, y_top,
- shape, rho, x_left=None, y_bottom=None):
+ def change_values(self, start_position, end_position, x_right=None,
+ y_top=None, shape=None, rho=None, x_left=None,
+ y_bottom=None, sym=True):
"""
Change the physical parameters between start_position and end_position.
@@ -563,41 +564,46 @@ class PhysicalModel:
----------
start_position : float
end_position : float
- x_right : float
+ x_right : float, optional
Right horizontal aperture in [m].
- y_top : float
+ y_top : float, optional
Top vertical aperture in [m].
- shape : str
+ shape : str, optional
Shape of the chamber cross section (elli/circ/rect).
- rho : float
+ rho : float, optional
Resistivity of the chamber material [ohm.m].
x_left : float, optional
Left horizontal aperture in [m].
y_bottom : float, optional
Bottom vertical aperture in [m].
+ sym : bool, optional
+ If True, right/left and top/bottum symmetry is applied.
"""
ind = (self.position > start_position) & (self.position < end_position)
- self.x_right[ind] = x_right
- self.y_top[ind] = y_top
-
- if x_left is None:
- self.x_left[ind] = -1*x_right
- else:
+ if x_right is not None:
+ self.x_right[ind] = x_right
+ if y_top is not None:
+ self.y_top[ind] = y_top
+ if x_left is not None:
self.x_left[ind] = x_left
-
- if y_bottom is None:
- self.y_bottom[ind] = -1*y_top
- else:
+ elif (x_right is not None) and (sym is True):
+ self.x_left[ind] = -1*x_right
+ if y_bottom is not None:
self.y_bottom[ind] = y_bottom
+ elif (y_top is not None) and (sym is True):
+ self.y_bottom[ind] = -1*y_top
ind2 = ((self.position[:-1] > start_position) &
(self.position[1:] < end_position))
- self.rho[ind2] = rho
- self.shape[ind2] = shape
-
- def taper(self, start_position, end_position, x_right_start, x_right_end,
- y_top_start, y_top_end, shape, rho, x_left_start=None,
- x_left_end=None, y_bottom_start=None, y_bottom_end=None):
+ if rho is not None:
+ self.rho[ind2] = rho
+ if shape is not None:
+ self.shape[ind2] = shape
+
+ def taper(self, start_position, end_position, x_right_start=None,
+ x_right_end=None, y_top_start=None, y_top_end=None, shape=None,
+ rho=None, x_left_start=None, x_left_end=None,
+ y_bottom_start=None, y_bottom_end=None, sym=True):
"""
Change the physical parameters to have a tapered transition between
start_position and end_position.
@@ -626,28 +632,31 @@ class PhysicalModel:
Bottom vertical aperture at the start of the taper in [m].
y_bottom_end : float, optional
Bottom vertical aperture at the end of the taper in [m].
+ sym : bool, optional
+ If True, right/left and top/bottum symmetry is applied.
"""
ind = (self.position > start_position) & (self.position < end_position)
- self.x_right[ind] = np.linspace(x_right_start, x_right_end, sum(ind))
- self.y_top[ind] = np.linspace(y_top_start, y_top_end, sum(ind))
-
+ if (x_right_start is not None) and (x_right_end is not None):
+ self.x_right[ind] = np.linspace(x_right_start, x_right_end, sum(ind))
+ if sym is True:
+ self.x_left[ind] = -1*np.linspace(x_right_start, x_right_end, sum(ind))
+
+ if (y_top_start is not None) and (y_top_end is not None):
+ self.y_top[ind] = np.linspace(y_top_start, y_top_end, sum(ind))
+ if sym is True:
+ self.y_bottom[ind] = -1*np.linspace(y_top_start, y_top_end, sum(ind))
+
if (x_left_start is not None) and (x_left_end is not None):
self.x_left[ind] = np.linspace(x_left_start, x_left_end, sum(ind))
- else:
- self.x_left[ind] = -1*np.linspace(x_right_start, x_right_end,
- sum(ind))
-
if (y_bottom_start is not None) and (y_bottom_end is not None):
- self.y_bottom[ind] = np.linspace(y_bottom_start, y_bottom_end,
- sum(ind))
- else:
- self.y_bottom[ind] = -1*np.linspace(y_top_start, y_top_end,
- sum(ind))
+ self.y_bottom[ind] = np.linspace(y_bottom_start, y_bottom_end, sum(ind))
- ind2 = ((self.position[:-1] > start_position)
- & (self.position[1:] < end_position))
- self.rho[ind2] = rho
- self.shape[ind2] = shape
+ ind2 = ((self.position[:-1] > start_position) &
+ (self.position[1:] < end_position))
+ if rho is not None:
+ self.rho[ind2] = rho
+ if shape is not None:
+ self.shape[ind2] = shape
def plot_aperture(self):
"""Plot horizontal and vertical apertures."""
@@ -664,3 +673,15 @@ class PhysicalModel:
ylabel="Vertical aperture [mm]")
axs[1].legend(["Top","Bottom"])
+ return (fig, axs)
+
+ def get_aperture(self, s):
+ self.xp = interp1d(self.position, self.x_right, kind='linear')
+ self.xm = interp1d(self.position, self.x_left, kind='linear')
+ self.yp = interp1d(self.position, self.y_top, kind='linear')
+ self.ym = interp1d(self.position, self.y_bottom, kind='linear')
+ aperture = np.array([self.xp(s),
+ self.xm(s),
+ self.yp(s),
+ self.ym(s)])
+ return aperture
\ No newline at end of file