Simple tracking test w/o MPI

Machines/soleil.py

@@ -26,15 +26,19 @@ def soleil(mode = 'Uniform'):
     emit = np.array([4.5e-9, 4.5e-9*0.01])
     sigma_0 = 8e-12
     sigma_delta = 9e-4
+    chro = [2,3]
     
     # mean values
     beta = np.array([3, 1.3])
     alpha = np.array([0, 0])
-    mean_val = Optics(beta, alpha)
+    disp = np.array([0, 0])
+    dispp = np.array([0, 0])
+    mean_val = Optics(beta, alpha, disp, dispp)
     
     ring = Synchrotron(h, L, E0, particle, ac=ac, U0=U0, tau=tau,
                        mean_optics=mean_val, emit=emit, tune=tune,
-                       sigma_delta=sigma_delta, sigma_0=sigma_0)
+                       sigma_delta=sigma_delta, sigma_0=sigma_0,
+                       chro=chro)
     
     if mode == 'Uniform':
         beam = Beam(ring, )

Tracking/one_turn_matrix.py

+# -*- coding: utf-8 -*-
+"""
+One turn map matrix elements
+
+@author: gamelina
+@date: 04/03/2020
+"""
+
+from abc import ABCMeta, abstractmethod
+from functools import wraps
+import numpy as np
+
+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.
+        """
+        raise NotImplementedError
+        
+    @staticmethod
+    def all_bunches(function):
+        """
+        """
+        @wraps(function)
+        def wrapper(*args, **kwargs):
+            for bunch in args[0]:
+                function(bunch, *args[1:], **kwargs)
+        return wrapper
+    
+    @staticmethod
+    def not_empty(function):
+        """
+        """
+        @wraps(function)
+        def wrapper(*args, **kwargs):
+            for bunch in args[0].not_empty:
+                function(bunch, *args[1:], **kwargs)
+        return wrapper
+
+        
+class Long_one_turn(Element):
+    """
+    Longitudinal transform for a single turn.
+    """
+    
+    def __init__(self, ring):
+        self.ring = ring
+        
+    def track_bunch(self, bunch):
+        """track"""
+        bunch["delta"] -= self.ring.U0/self.ring.E0
+        bunch["tau"] -= self.ring.ac*self.ring.T0*bunch["delta"]
+        
+    def track(self, beam):
+        """track full beam"""
+        super().not_empty(self.track_bunch)(beam)
+        
+class SynchrotronRadiation(Element):
+    """SyncRad"""
+    
+    def __init__(self, ring):
+        self.ring = ring
+        
+    def track_bunch(self, bunch):
+        """track"""
+        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)
+        
+    def track(self, beam):
+        """track full beam"""
+        super().not_empty(self.track_bunch)(beam)
+
+class RF_cavity(Element):
+    """ Perfect RF cavity class for main and harmonic RF cavities."""
+    def __init__(self, ring, m, Vc, theta):
+        self.ring = ring
+        self.m = m # Harmonic number of the cavity
+        self.Vc = Vc # Amplitude of Cavity voltage [V]
+        self.theta = theta # phase of Cavity voltage
+        
+    def track_bunch(self,bunch):
+        """Track """
+        bunch["delta"] += self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*bunch["tau"] + self.theta)
+        
+    def track(self, beam):
+        """track full beam"""
+        super().not_empty(self.track_bunch)(beam)
+        
+class Trans_one_turn(Element):
+    """
+    Transverse transformation for a single turn.
+    """
+    
+    def __init__(self, ring):
+        self.ring = ring
+        self.alpha = self.ring.mean_optics.alpha
+        self.beta = self.ring.mean_optics.beta
+        self.gamma = self.ring.mean_optics.gamma
+        self.disp = self.ring.mean_optics.disp
+        self.dispp = self.ring.mean_optics.dispp
+        self.phase_advance = self.ring.tune[0:2]*2*np.pi
+        
+    def track_bunch(self, bunch):
+        """track"""
+
+        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"])
+        matrix = np.zeros((6,6,len(bunch)))
+        matrix[0,0,:] = np.cos(phase_advance_x) + self.alpha[0]*np.sin(phase_advance_x)
+        matrix[0,1,:] = self.beta[0]*np.cos(phase_advance_x)
+        matrix[0,2,:] = self.disp[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.dispp[0]
+        matrix[2,2,:] = 1
+        
+        matrix[3,3,:] = np.cos(phase_advance_y) + self.alpha[1]*np.sin(phase_advance_y)
+        matrix[3,4,:] = self.beta[1]*np.cos(phase_advance_y)
+        matrix[3,5,:] = self.disp[1]
+        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.dispp[1]
+        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 track(self, beam):
+        """track full beam"""
+        super().not_empty(self.track_bunch)(beam)

Tracking/synchrotron.py

@@ -41,9 +41,11 @@ class Proton(Particle):
 
 class Optics:
     """Handle optic functions"""
-    def __init__(self, beta, alpha):
+    def __init__(self, beta, alpha, disp, dispp):
         self.beta = beta
         self.alpha = alpha
+        self.disp = disp
+        self.dispp = dispp
 
     @property
     def gamma(self):