Select Git revision
-
Stephane Poirier authoredStephane Poirier authored
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
element.py 7.39 KiB
# -*- 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.
@author: gamelina
@date: 11/03/2020
"""
import numpy as np
from abc import ABCMeta, abstractmethod
from functools import wraps
from tracking.particles import Beam
class Element(metaclass=ABCMeta):
"""
Abstract Element class used for subclass inheritance to define all kinds
of objects which intervene in the tracking.
"""
@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
@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
class LongitudinalMap(Element):
"""
Longitudinal map for a single turn in the synchrotron.
Parameters
----------
ring : Synchrotron object
"""
def __init__(self, ring):
self.ring = ring
@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.ac * self.ring.T0 * bunch["delta"]
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
@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),2))
bunch["x"] += self.ring.sigma[0]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,0]
bunch["xp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["xp"]
bunch["xp"] += self.ring.sigma[1]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,1]
if (self.switch[2] == True):
rand = np.random.normal(size=(len(bunch),2))
bunch["y"] += self.ring.sigma[2]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,0]
bunch["yp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["yp"]
bunch["yp"] += self.ring.sigma[3]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,1]
# Reset energy change to 0 for next turn
bunch.energy_change = 0
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
self.phase_advance = self.ring.tune[0:2]*2*np.pi
@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 adcence which depends on energy via chromaticity
phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])
phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
matrix = np.zeros((6,6,len(bunch)))
# Horizontal
matrix[0,0,:] = np.cos(phase_advance_x) + self.alpha[0]*np.sin(phase_advance_x)
matrix[0,1,:] = self.beta[0]*np.sin(phase_advance_x)
matrix[0,2,:] = self.dispersion[0]
matrix[1,0,:] = -1*self.gamma[0]*np.sin(phase_advance_x)
matrix[1,1,:] = np.cos(phase_advance_x) - self.alpha[0]*np.sin(phase_advance_x)
matrix[1,2,:] = self.dispersion[1]
matrix[2,2,:] = 1
# Vertical
matrix[3,3,:] = np.cos(phase_advance_y) + self.alpha[1]*np.sin(phase_advance_y)
matrix[3,4,:] = self.beta[1]*np.sin(phase_advance_y)
matrix[3,5,:] = self.dispersion[2]
matrix[4,3,:] = -1*self.gamma[1]*np.sin(phase_advance_y)
matrix[4,4,:] = np.cos(phase_advance_y) - self.alpha[1]*np.sin(phase_advance_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