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

Merge branch 'ImpedanceModel' into develop

parents e0b23f39 0a0da640
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
+ 97
− 135
View file @ 06f396b1
......@@ -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()
mbtrack2/utilities/optics.py
+ 59
− 38
View file @ 06f396b1
......@@ -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
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