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

Add TransverseMapSector element 

Update the Optics class to be able to compute the phase advance at any location on the lattice.
Add the TransverseMapSector element to be able to split the lattice in several parts.
Add the transverse_map_sector_generator convenience function to split a lattice in several TransverseMapSector elements.
parent 58370b5b
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
mbtrack2/tracking/__init__.py
+ 3
− 1
View file @ d8075bf9
......@@ -12,7 +12,9 @@ from mbtrack2.tracking.element import (Element,
LongitudinalMap,
TransverseMap,
SynchrotronRadiation,
SkewQuadrupole)
SkewQuadrupole,
TransverseMapSector,
transverse_map_sector_generator)
from mbtrack2.tracking.aperture import (CircularAperture,
ElipticalAperture,
RectangularAperture,
......
mbtrack2/tracking/element.py
+ 158
− 0
View file @ d8075bf9
......@@ -243,3 +243,161 @@ class SkewQuadrupole:
bunch['xp'] = bunch['xp'] - self.strength * bunch['y']
bunch['yp'] = bunch['yp'] - self.strength * bunch['x']
class TransverseMapSector(Element):
"""
Transverse map for a sector of the synchrotron, from an initial
position s0 to a final position s1.
Parameters
----------
ring : Synchrotron object
Ring parameters.
alpha0 : array of shape (2,)
Alpha Twiss function at the initial location of the sector.
beta0 : array of shape (2,)
Beta Twiss function at the initial location of the sector.
dispersion0 : array of shape (4,)
Dispersion function at the initial location of the sector.
alpha1 : array of shape (2,)
Alpha Twiss function at the final location of the sector.
beta1 : array of shape (2,)
Beta Twiss function at the final location of the sector.
dispersion1 : array of shape (4,)
Dispersion function at the final location of the sector.
phase_diff : array of shape (2,)
Phase difference between the initial and final location of the
sector.
chro_diff : array of shape (2,)
Chromaticity difference between the initial and final location of
the sector.
adts : array of shape (4,), optional
Amplitude-dependent tune shift of the sector, see Synchrotron class
for details. The default is None.
"""
def __init__(self, ring, alpha0, beta0, dispersion0, alpha1, beta1,
dispersion1, phase_diff, chro_diff, adts=None):
self.ring = ring
self.alpha0 = alpha0
self.beta0 = beta0
self.dispersion0 = dispersion0
self.alpha1 = alpha1
self.beta1 = beta1
self.dispersion1 = dispersion1
self.tune_diff = phase_diff / (2*np.pi)
self.chro_diff = chro_diff
if adts is not None:
self.adts_poly = [np.poly1d(adts[0]),
np.poly1d(adts[1]),
np.poly1d(adts[2]),
np.poly1d(adts[3])]
else:
self.adts_poly = None
@Element.parallel
def track(self, bunch):
"""
Tracking method for the element.
No bunch to bunch interaction, so written for Bunch objects and
@Element.parallel is used to handle Beam objects.
Parameters
----------
bunch : Bunch or Beam object
"""
# Compute phase advance which depends on energy via chromaticity and ADTS
if self.adts_poly is None:
phase_advance_x = 2*np.pi * (self.tune_diff[0] +
self.chro_diff[0]*bunch["delta"])
phase_advance_y = 2*np.pi * (self.tune_diff[1] +
self.chro_diff[1]*bunch["delta"])
else:
phase_advance_x = 2*np.pi * (self.tune_diff[0] +
self.chro_diff[0]*bunch["delta"] +
self.adts_poly[0](bunch['x']) +
self.adts_poly[2](bunch['y']))
phase_advance_y = 2*np.pi * (self.tune_diff[1] +
self.chro_diff[1]*bunch["delta"] +
self.adts_poly[1](bunch['x']) +
self.adts_poly[3](bunch['y']))
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
matrix = np.zeros((6,6,len(bunch)))
# Horizontal
matrix[0,0,:] = np.sqrt(self.beta1[0]/self.beta0[0])*(np.cos(phase_advance_x) + self.alpha0[0]*np.sin(phase_advance_x))
matrix[0,1,:] = np.sqrt(self.beta0[0]*self.beta1[0])*np.sin(phase_advance_x)
matrix[0,2,:] = self.dispersion1[0] - matrix[0,0,:]*self.dispersion0[0] - matrix[0,1,:]*self.dispersion0[1]
matrix[1,0,:] = ((self.alpha0[0] - self.alpha1[0])*np.cos(phase_advance_x) - (1 + self.alpha0[0]*self.alpha1[0])*np.sin(phase_advance_x))/np.sqrt(self.beta0[0]*self.beta1[0])
matrix[1,1,:] = np.sqrt(self.beta0[0]/self.beta1[0])*(np.cos(phase_advance_x) - self.alpha1[0]*np.sin(phase_advance_x))
matrix[1,2,:] = self.dispersion1[1] - matrix[1,0,:]*self.dispersion0[0] - matrix[1,1,:]*self.dispersion0[1]
matrix[2,2,:] = 1
# Vertical
matrix[3,3,:] = np.sqrt(self.beta1[1]/self.beta0[1])*(np.cos(phase_advance_y) + self.alpha0[1]*np.sin(phase_advance_y))
matrix[3,4,:] = np.sqrt(self.beta0[1]*self.beta1[1])*np.sin(phase_advance_y)
matrix[3,5,:] = self.dispersion1[2] - matrix[3,3,:]*self.dispersion0[2] - matrix[3,4,:]*self.dispersion0[3]
matrix[4,3,:] = ((self.alpha0[1] - self.alpha1[1])*np.cos(phase_advance_y) - (1 + self.alpha0[1]*self.alpha1[1])*np.sin(phase_advance_y))/np.sqrt(self.beta0[1]*self.beta1[1])
matrix[4,4,:] = np.sqrt(self.beta0[1]/self.beta1[1])*(np.cos(phase_advance_y) - self.alpha1[1]*np.sin(phase_advance_y))
matrix[4,5,:] = self.dispersion1[3] - matrix[4,3,:]*self.dispersion0[2] - matrix[4,4,:]*self.dispersion0[3]
matrix[5,5,:] = 1
x = matrix[0,0,:]*bunch["x"] + matrix[0,1,:]*bunch["xp"] + matrix[0,2,:]*bunch["delta"]
xp = matrix[1,0,:]*bunch["x"] + matrix[1,1,:]*bunch["xp"] + matrix[1,2,:]*bunch["delta"]
y = matrix[3,3,:]*bunch["y"] + matrix[3,4,:]*bunch["yp"] + matrix[3,5,:]*bunch["delta"]
yp = matrix[4,3,:]*bunch["y"] + matrix[4,4,:]*bunch["yp"] + matrix[4,5,:]*bunch["delta"]
bunch["x"] = x
bunch["xp"] = xp
bunch["y"] = y
bunch["yp"] = yp
def transverse_map_sector_generator(ring, positions):
"""
Convenience function which generate a list of TransverseMapSector elements
from a ring with an AT lattice.
Tracking through all the sectors is equivalent to a full turn (and thus to
the TransverseMap object).
Parameters
----------
ring : Synchrotron object
Ring parameters, must .
positions : array
List of longitudinal positions in [m] to use as starting and end points
of the TransverseMapSector elements.
Returns
-------
sectors : list
List of TransverseMapSector elements.
"""
if ring.optics.use_local_values:
raise ValueError("The Synchrotron object must be loaded from an AT lattice")
N_sec = len(positions)
sectors = []
chro_diff = ring.chro/N_sec
for i in range(N_sec):
alpha0 = ring.optics.alpha(positions[i])
beta0 = ring.optics.beta(positions[i])
dispersion0 = ring.optics.dispersion(positions[i])
mu0 = ring.optics.mu(positions[i])
if i != (N_sec - 1):
alpha1 = ring.optics.alpha(positions[i+1])
beta1 = ring.optics.beta(positions[i+1])
dispersion1 = ring.optics.dispersion(positions[i+1])
mu1 = ring.optics.mu(positions[i+1])
else:
alpha1 = ring.optics.alpha(positions[0])
beta1 = ring.optics.beta(positions[0])
dispersion1 = ring.optics.dispersion(positions[0])
mu1 = ring.optics.mu(ring.L)
phase_diff = mu1 - mu0
sectors.append(TransverseMapSector(ring, alpha0, beta0, dispersion0,
alpha1, beta1, dispersion1, phase_diff, chro_diff))
return sectors
mbtrack2/utilities/optics.py
+ 36
− 5
View file @ d8075bf9
......@@ -123,6 +123,7 @@ class Optics:
self.beta_array = np.tile(twiss.beta.T, self.periodicity)
self.alpha_array = np.tile(twiss.alpha.T, self.periodicity)
self.dispersion_array = np.tile(twiss.dispersion.T, self.periodicity)
self.mu_array = np.tile(twiss.mu.T, self.periodicity)
self.position = np.append(self.position, self.lattice.circumference)
self.beta_array = np.append(self.beta_array, self.beta_array[:,0:1],
......@@ -131,12 +132,17 @@ class Optics:
axis=1)
self.dispersion_array = np.append(self.dispersion_array,
self.dispersion_array[:,0:1], axis=1)
self.mu_array = np.append(self.mu_array,
self.mu_array[:,0:1], axis=1)
self.gamma_array = (1 + self.alpha_array**2)/self.beta_array
self.tune = tune * self.periodicity
self.chro = chrom * self.periodicity
self.ac = at.get_mcf(self.lattice)
self.mu_array[:,-1] = (np.floor(self.mu_array[:,-2]/(2*np.pi)) +
self.tune)*2*np.pi
self.setup_interpolation()
......@@ -162,6 +168,10 @@ class Optics:
kind='linear')
self.disppY = interp1d(self.position, self.dispersion_array[3,:],
kind='linear')
self.muX = interp1d(self.position, self.mu_array[0,:],
kind='linear')
self.muY = interp1d(self.position, self.mu_array[1,:],
kind='linear')
@property
def local_beta(self):
......@@ -303,13 +313,34 @@ class Optics:
self.dispY(position), self.disppY(position)]
return np.array(dispersion)
def mu(self, position):
"""
Return phase advances at specific locations given by position.
If no lattice has been loaded, None is returned.
Parameters
----------
position : array or float
Longitudinal position at which the phase advances are returned.
Returns
-------
mu : array
Phase advances.
"""
if self.use_local_values:
return np.outer(np.array([0,0]), np.ones_like(position))
else:
mu = [self.muX(position), self.muY(position)]
return np.array(mu)
def plot(self, var, option, n_points=1000):
"""
Plot optical variables.
Parameters
----------
var : {"beta", "alpha", "gamma", "dispersion"}
var : {"beta", "alpha", "gamma", "dispersion", "mu"}
Optical variable to be plotted.
option : str
If var = "beta", "alpha" and "gamma", option = {"x","y"} specifying
......@@ -322,7 +353,7 @@ class Optics:
"""
var_dict = {"beta":self.beta, "alpha":self.alpha, "gamma":self.gamma,
"dispersion":self.dispersion}
"dispersion":self.dispersion, "mu":self.mu}
if var == "dispersion":
option_dict = {"x":0, "px":1, "y":2, "py":3}
......@@ -332,15 +363,15 @@ class Optics:
ylabel = label[option_dict[option]]
elif var=="beta" or var=="alpha" or var=="gamma":
elif var=="beta" or var=="alpha" or var=="gamma" or var=="mu":
option_dict = {"x":0, "y":1}
label_dict = {"beta":"$\\beta$", "alpha":"$\\alpha$",
"gamma":"$\\gamma$"}
"gamma":"$\\gamma$", "mu":"$\\mu$"}
if option == "x": label_sup = "$_{x}$"
elif option == "y": label_sup = "$_{y}$"
unit = {"beta":" (m)", "alpha":"", "gamma":" (m$^{-1}$)"}
unit = {"beta":" (m)", "alpha":"", "gamma":" (m$^{-1}$)", "mu":""}
ylabel = label_dict[var] + label_sup + unit[var]
......
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