Radiation damping in transverse plane and bug fix

Add synchrotron radiation in transverse plane
Fix a mistake in the transverse one turn map

Tracking/beam.py

@@ -64,6 +64,8 @@ class Bunch:
         if not alive:
             self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool))  
             
+        self._energy_change = pd.Series(np.zeros((self.mp_number,),dtype=float))  
+        
     def __len__(self):
         """Return the number of alive particles"""
         return len(self[:])
@@ -84,6 +86,18 @@ class Bunch:
         """Return representation of alive particles"""
         return f'{self[:]!r}'
     
+    @property
+    def energy_change(self):
+        """Store the particle relative energy change in the last turn. Used by
+        the SynchrotronRadiation class to compute radiation damping. To be
+        changed by Element objects which change the energy, for example RF
+        cavities."""
+        return self._energy_change.loc[self.alive]
+    
+    @energy_change.setter
+    def energy_change(self, value):
+        self._energy_change.loc[self.alive] = value 
+        
     @property
     def mp_number(self):
         """Macro-particle number"""

Tracking/one_turn_matrix.py

@@ -72,15 +72,28 @@ class Long_one_turn(Element):
 class SynchrotronRadiation(Element):
     """SyncRad"""
     
-    def __init__(self, ring):
+    def __init__(self, ring, switch = np.ones((3,), dtype=bool)):
         self.ring = ring
+        self.switch = switch
         
     @Element.not_empty        
     def track(self, bunch):
         """track"""
+        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]
+        bunch.energy_change = 0
 
 class RF_cavity(Element):
     """ Perfect RF cavity class for main and harmonic RF cavities."""
@@ -93,8 +106,12 @@ class RF_cavity(Element):
     @Element.not_empty    
     def track(self,bunch):
         """Track """
-        bunch["delta"] += self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*bunch["tau"] + self.theta)
+        energy_change = self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*bunch["tau"] + self.theta)
+        bunch["delta"] += energy_change
+        bunch.energy_change += energy_change
         
+    def value(self, val):
+        return self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*val + self.theta)
         
 class Trans_one_turn(Element):
     """
@@ -118,7 +135,7 @@ class Trans_one_turn(Element):
         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,1,:] = self.beta[0]*np.sin(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)
@@ -126,7 +143,7 @@ class Trans_one_turn(Element):
         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,4,:] = self.beta[1]*np.sin(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)

Tracking/parallel.py

@@ -19,8 +19,7 @@ class Mpi:
             from mpi4py import MPI
         except(ModuleNotFoundError):
             print("Please install mpi4py module.")
-        else:
-            print("MPI error.")
+
         self.comm = MPI.COMM_WORLD
         self.rank = self.comm.Get_rank()
         self.size = self.comm.Get_size()

Tracking/synchrotron.py

@@ -182,3 +182,12 @@ class Synchrotron:
         """Momentum compaction"""
         return self.ac - 1/(self.gamma**2)
     
+    @property
+    def sigma(self):
+        """RMS beam size at equilibrium"""
+        sigma = np.zeros((4,))
+        sigma[0] = (self.emit[0]*self.mean_optics.beta[0] + self.mean_optics.disp[0]**2*self.sigma_delta)**0.5
+        sigma[1] = (self.emit[0]*self.mean_optics.alpha[0] + self.mean_optics.dispp[0]**2*self.sigma_delta)**0.5
+        sigma[2] = (self.emit[1]*self.mean_optics.beta[1] + self.mean_optics.disp[1]**2*self.sigma_delta)**0.5
+        sigma[3] = (self.emit[1]*self.mean_optics.alpha[1] + self.mean_optics.dispp[1]**2*self.sigma_delta)**0.5
+        return sigma
\ No newline at end of file