Skip to content
Snippets Groups Projects
Select Git revision
  • main
  • hdf5_1_8_x
  • hdf5_1_14_3
  • compression
  • version-3_4_x
  • NexusCPP-4.0.2
  • NexusCPP-3.5.4
  • NexusCPP-3.5.3
  • NexusCPP-3.5.2
  • NexusCPP-3.5.1
  • NexusCPP-3.5.0
  • NexusCPP-3.4.1
  • NexusCPP-3.4.0
  • NexusCPP-3.3.7
  • NexusCPP-3.3.6
  • NexusCPP-3.3.5
  • NexusCPP-3.3.4
  • NexusCPP-3.3.3
  • NexusCPP-3.3.2
  • NexusCPP-3.3.1
  • NexusCPP-3.3.0
  • NexusCPP-3.2.2
  • NexusCPP-3.2.1
  • NexusCPP-3.2.0
  • NexusCPP-3.1.6
25 results

README

Blame
  • 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