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
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
mbtrack2/impedance/impedance_model.py
+ 92
− 131
View file @ 641a667c
......@@ -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):
"""
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

The place to be for your source code