Skip to content
Snippets Groups Projects
Commit 641a667c authored by Alexis GAMELIN's avatar Alexis GAMELIN
Browse files

Rework ImpedanceModel 

Simplify and optimize adding element into ImpedanceModel.
parent 58370b5b
Branches
Tags
No related merge requests found
Show whitespace changes
Inline Side-by-side
mbtrack2/impedance/impedance_model.py
+ 92
− 131
View file @ 641a667c
...@@ -6,6 +6,7 @@ import pandas as pd ...@@ -6,6 +6,7 @@ import pandas as pd
import numpy as np import numpy as np
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import pickle import pickle
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
...@@ -14,52 +15,43 @@ from mbtrack2.utilities.misc import (beam_loss_factor, effective_impedance, ...@@ -14,52 +15,43 @@ from mbtrack2.utilities.misc import (beam_loss_factor, effective_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
from mbtrack2.tracking.element import Element
class ImpedanceModel(Element): class ImpedanceModel():
""" """
Define the impedance model of the machine. Define the impedance model of the machine.
The model must be completed with successive add(...) calls, then
compute_sum() must be run.
Parameters Parameters
---------- ----------
ring : Synchrotron object 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 Attributes
---------- ----------
wakefields : list of WakeField objects wakefields : list of WakeField objects
WakeFields in the model. WakeFields in the model.
positions : array positions : list of arrays
WakeFields positions. Positions corresponding the different WakeFields.
names : array names : list of str
Names of (unique) WakeField objects. Names of the WakeField objects.
sum : WakeField sum : WakeField
Sum of every WakeField in the model. Sum of every WakeField in the model weigthed by beta functions.
sum_"name" : WakeField sum_"name" : WakeField
Sum of every Wakefield with the same "name". 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.
Methods Methods
------- -------
add(wakefield_list, wakefiled_positions) add(wakefield, positions, name)
Add elements to the model. Add the same WakeField object at different locations to the model.
add_multiple_elements(wakefield, wakefiled_positions) sum_beta(wake, beta)
Add the same element at different locations to the model. Weight a WakeField object by an array of beta functions.
sum_elements() compute_sum_names()
Sum all WakeFields into self.sum. Compute the weighted WakeField for each WakeField object type.
update_name_list() compute_sum()
Update self.names with uniques names of self.wakefields. Compute the sum of all weighted WakeField into self.sum.
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.
plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None) plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
Plot the contributions of different kind of WakeFields. Plot the contributions of different kind of WakeFields.
save(file) save(file)
...@@ -68,131 +60,100 @@ class ImpedanceModel(Element): ...@@ -68,131 +60,100 @@ class ImpedanceModel(Element):
Load impedance model from file. Load impedance model from file.
""" """
def __init__(self, ring, wakefield_list=None, wakefiled_positions=None): def __init__(self, ring):
self.ring = ring self.ring = ring
self.optics = self.ring.optics self.optics = self.ring.optics
self.wakefields = [] self.wakefields = []
self.positions = np.array([]) self.positions = []
self.names = np.array([]) self.names = []
self.sum_names = np.array([]) self.sum_names = []
self.add(wakefield_list, wakefiled_positions)
def track(self, beam):
"""
Track a beam object through this Element.
Parameters def add(self, wakefield, positions, name=None):
----------
beam : Beam object
"""
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):
""" """
Return indexes of WakeFields with the same name in self.wakefields. Add the same WakeField object at different locations to the model.
Parameters Parameters
---------- ----------
name : str wakefield : WakeField
WakeField name. 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 Returns
------- -------
index : list None.
Index of positions in self.wakefields.
""" """
index = [] self.wakefields.append(wakefield)
for i, wakefield in enumerate(self.wakefields): self.positions.append(positions)
if wakefield.name == name: if name is None:
index.append(i) name = wakefield.name
return index 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.")
def sum_by_name(self, name): @staticmethod
def sum_beta(wake, beta):
""" """
Sum the elements with the same name in the model into sum_name. Weight a WakeField object by an array of beta functions.
Parameters Parameters
---------- ----------
name : str wake : WakeField
Name of the WakeField to sum. WakeField element object.
""" beta : array of shape (2, N)
attribute_name = "sum_" + name Beta function at the locations of the elements.
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): Returns
""" -------
Sum all the elements with the same name in the model into sum_name. wake_sum : WakeField
""" WakeField object weighted by beta functions.
for name in self.names:
self.sum_by_name(name)
def add(self, wakefield_list, wakefiled_positions):
""" """
Add elements to the model. wake_sum = deepcopy(wake)
for component_name in wake.components:
Parameters comp = getattr(wake_sum, component_name)
---------- weight = ((beta[0,:] ** comp.power_x) *
wakefield_list : list of WakeField objects (beta[1,:] ** comp.power_y))
WakeFields to add to the model. setattr(wake_sum, component_name, weight.sum()*comp)
wakefiled_positions : list return wake_sum
Longitudinal positions corresponding to the added Wakfields.
def compute_sum_names(self):
""" """
if (wakefield_list is not None) and (wakefiled_positions is not None): Compute the weighted WakeField for each WakeField object type.
for wakefield in wakefield_list: The new summed WakeField object are set to into self.sum_name.
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):
""" """
Add the same element at different locations to the model. 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)
Parameters def compute_sum(self):
----------
WakeField : WakeField object
WakeField to add to the model.
wakefiled_positions : list
Longitudinal positions corresponding to the added Wakfield.
""" """
for position in wakefiled_positions: Compute the sum of all weighted WakeField into self.sum.
self.positions = np.append(self.positions, position) self.compute_sum_names must be called before this.
self.wakefields.append(wakefield) """
self.compute_sum_names()
self.update_name_list() 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)
def plot_area(self, Z_type="Zlong", component="real", sigma=None, def plot_area(self, Z_type="Zlong", component="real", sigma=None,
attr_list=None, zoom=False): attr_list=None, zoom=False):
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment

The place to be for your source code