# -*- coding: utf-8 -*-
"""
This module defines the most basic elements for tracking, including Element,
an abstract base class which is to be used as mother class to every elements
included in the tracking.
"""
from abc import ABCMeta, abstractmethod
from copy import deepcopy
from functools import wraps
import numpy as np
from mbtrack2.tracking.particles import Beam
[docs]class Element(metaclass=ABCMeta):
"""
Abstract Element class used for subclass inheritance to define all kinds
of objects which intervene in the tracking.
"""
[docs] @abstractmethod
def track(self, beam):
"""
Track a beam object through this Element.
This method needs to be overloaded in each Element subclass.
Parameters
----------
beam : Beam object
"""
raise NotImplementedError
[docs] @staticmethod
def parallel(track):
"""
Defines the decorator @parallel which handle the embarrassingly
parallel case which happens when there is no bunch to bunch
interaction in the tracking routine.
Adding @Element.parallel allows to write the track method of the
Element subclass for a Bunch object instead of a Beam object.
Parameters
----------
track : function, method of an Element subclass
track method of an Element subclass which takes a Bunch object as
input
Returns
-------
track_wrapper: function, method of an Element subclass
track method of an Element subclass which takes a Beam object or a
Bunch object as input
"""
@wraps(track)
def track_wrapper(*args, **kwargs):
if isinstance(args[1], Beam):
self = args[0]
beam = args[1]
if (beam.mpi_switch == True):
track(self, beam[beam.mpi.bunch_num], *args[2:], **kwargs)
else:
for bunch in beam.not_empty:
track(self, bunch, *args[2:], **kwargs)
else:
self = args[0]
bunch = args[1]
track(self, bunch, *args[2:], **kwargs)
return track_wrapper
[docs]class LongitudinalMap(Element):
"""
Longitudinal map for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
"""
def __init__(self, ring):
self.ring = ring
[docs] @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
"""
bunch["delta"] -= self.ring.U0 / self.ring.E0
bunch["tau"] += self.ring.eta(
bunch["delta"]) * self.ring.T0 * bunch["delta"]
[docs]class SynchrotronRadiation(Element):
"""
Element to handle synchrotron radiation, radiation damping and quantum
excitation, for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
switch : bool array of shape (3,), optional
allow to choose on which plane the synchrotron radiation is active
"""
def __init__(self, ring, switch=np.ones((3, ), dtype=bool)):
self.ring = ring
self.switch = switch
[docs] @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
"""
if (self.switch[0] == True):
rand = np.random.normal(size=len(bunch))
bunch["delta"] = (
(1 - 2 * self.ring.T0 / self.ring.tau[2]) * bunch["delta"] +
2 * self.ring.sigma_delta *
(self.ring.T0 / self.ring.tau[2])**0.5 * rand)
if (self.switch[1] == True):
rand = np.random.normal(size=len(bunch))
bunch["xp"] = (
(1 - 2 * self.ring.T0 / self.ring.tau[0]) * bunch["xp"] +
2 * self.ring.sigma()[1] *
(self.ring.T0 / self.ring.tau[0])**0.5 * rand)
if (self.switch[2] == True):
rand = np.random.normal(size=len(bunch))
bunch["yp"] = (
(1 - 2 * self.ring.T0 / self.ring.tau[1]) * bunch["yp"] +
2 * self.ring.sigma()[3] *
(self.ring.T0 / self.ring.tau[1])**0.5 * rand)
[docs]class TransverseMap(Element):
"""
Transverse map for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
"""
def __init__(self, ring):
self.ring = ring
self.alpha = self.ring.optics.local_alpha
self.beta = self.ring.optics.local_beta
self.gamma = self.ring.optics.local_gamma
self.dispersion = self.ring.optics.local_dispersion
if self.ring.adts is not None:
self.adts_poly = [
np.poly1d(self.ring.adts[0]),
np.poly1d(self.ring.adts[1]),
np.poly1d(self.ring.adts[2]),
np.poly1d(self.ring.adts[3])
]
[docs] @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.ring.adts is None:
phase_advance_x = 2 * np.pi * (self.ring.tune[0] +
self.ring.chro[0] * bunch["delta"])
phase_advance_y = 2 * np.pi * (self.ring.tune[1] +
self.ring.chro[1] * bunch["delta"])
else:
Jx = (self.ring.optics.local_gamma[0] * bunch['x']**2) + \
(2*self.ring.optics.local_alpha[0] * bunch['x']*bunch['xp']) + \
(self.ring.optics.local_beta[0] * bunch['xp']**2)
Jy = (self.ring.optics.local_gamma[1] * bunch['y']**2) + \
(2*self.ring.optics.local_alpha[1] * bunch['y']*bunch['yp']) + \
(self.ring.optics.local_beta[1] * bunch['yp']**2)
phase_advance_x = 2 * np.pi * (
self.ring.tune[0] + self.ring.chro[0] * bunch["delta"] +
self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
phase_advance_y = 2 * np.pi * (
self.ring.tune[1] + self.ring.chro[1] * bunch["delta"] +
self.adts_poly[1](Jx) + self.adts_poly[3](Jy))
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
matrix = np.zeros((6, 6, len(bunch)), dtype=np.float64)
# Horizontal
c_x = np.cos(phase_advance_x)
s_x = np.sin(phase_advance_x)
matrix[0, 0, :] = c_x + self.alpha[0] * s_x
matrix[0, 1, :] = self.beta[0] * s_x
matrix[0, 2, :] = self.dispersion[0]
matrix[1, 0, :] = -1 * self.gamma[0] * s_x
matrix[1, 1, :] = c_x - self.alpha[0] * s_x
matrix[1, 2, :] = self.dispersion[1]
matrix[2, 2, :] = 1
# Vertical
c_y = np.cos(phase_advance_y)
s_y = np.sin(phase_advance_y)
matrix[3, 3, :] = c_y + self.alpha[1] * s_y
matrix[3, 4, :] = self.beta[1] * s_y
matrix[3, 5, :] = self.dispersion[2]
matrix[4, 3, :] = -1 * self.gamma[1] * s_y
matrix[4, 4, :] = c_y - self.alpha[1] * s_y
matrix[4, 5, :] = self.dispersion[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
[docs]class SkewQuadrupole:
"""
Thin skew quadrupole element used to introduce betatron coupling (the
length of the quadrupole is neglected).
Parameters
----------
strength : float
Integrated strength of the skew quadrupole [m].
"""
def __init__(self, strength):
self.strength = strength
[docs] @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
"""
bunch['xp'] = bunch['xp'] - self.strength * bunch['y']
bunch['yp'] = bunch['yp'] - self.strength * bunch['x']
[docs]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.gamma0 = (1 + self.alpha0**2) / self.beta0
self.dispersion0 = dispersion0
self.alpha1 = alpha1
self.beta1 = beta1
self.gamma1 = (1 + self.alpha1**2) / self.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
[docs] @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:
Jx = (self.gamma0[0] * bunch['x']**2) + \
(2*self.alpha0[0] * bunch['x']*self['xp']) + \
(self.beta0[0] * bunch['xp']**2)
Jy = (self.gamma0[1] * bunch['y']**2) + \
(2*self.alpha0[1] * bunch['y']*bunch['yp']) + \
(self.beta0[1] * bunch['yp']**2)
phase_advance_x = 2 * np.pi * (
self.tune_diff[0] + self.chro_diff[0] * bunch["delta"] +
self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
phase_advance_y = 2 * np.pi * (
self.tune_diff[1] + self.chro_diff[1] * bunch["delta"] +
self.adts_poly[1](Jx) + self.adts_poly[3](Jy))
# 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
[docs]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.
The array should contain the initial position (s=0) but not the end
position (s=ring.L), so like position = np.array([0, pos1, pos2, ...]).
Returns
-------
sectors : list
List of TransverseMapSector elements.
"""
import at
def _compute_chro(ring, pos, dp=1e-4):
lat = deepcopy(ring.optics.lattice)
lat.append(at.Marker("END"))
N = len(lat)
refpts = np.arange(N)
*elem_neg_dp, = at.linopt2(lat, refpts=refpts, dp=-dp)
*elem_pos_dp, = at.linopt2(lat, refpts=refpts, dp=dp)
s = elem_neg_dp[2]["s_pos"]
mux0 = elem_neg_dp[2]['mu'][:, 0]
mux1 = elem_pos_dp[2]['mu'][:, 0]
muy0 = elem_neg_dp[2]['mu'][:, 1]
muy1 = elem_pos_dp[2]['mu'][:, 1]
Chrox = (mux1-mux0) / (2*dp) / 2 / np.pi
Chroy = (muy1-muy0) / (2*dp) / 2 / np.pi
chrox = np.interp(pos, s, Chrox)
chroy = np.interp(pos, s, Chroy)
return np.array([chrox, chroy])
if ring.optics.use_local_values:
raise ValueError(
"The Synchrotron object must be loaded from an AT lattice")
N_sec = len(positions)
sectors = []
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])
chro0 = _compute_chro(ring, 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])
chro1 = _compute_chro(ring, 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)
chro1 = _compute_chro(ring, ring.L)
phase_diff = mu1 - mu0
chro_diff = chro1 - chro0
sectors.append(
TransverseMapSector(ring, alpha0, beta0, dispersion0, alpha1,
beta1, dispersion1, phase_diff, chro_diff))
return sectors