Source code for mbtrack2.tracking.element

# -*- 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