Add Courant-Snyder invariant calculation method to Bunch class

tracking/element.py

@@ -142,6 +142,27 @@ class SynchrotronRadiation(Element):
         # Reset energy change to 0 for next turn
         bunch.energy_change = 0
         
+        #--------------------------------------
+        # 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.

tracking/monitors/__init__.py

@@ -15,4 +15,5 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
                                                  plot_phasespacedata,
                                                  plot_profiledata,
                                                  plot_beamdata,
-                                                 plot_wakedata)
+                                                 plot_wakedata,
\ No newline at end of file
+                                                 plot_tunedata)
\ No newline at end of file

tracking/monitors/monitors.py

@@ -16,7 +16,7 @@ from mbtrack2.tracking.element import Element
 from mbtrack2.tracking.particles import Bunch, Beam
 from scipy.interpolate import interp1d
 from abc import ABCMeta
-from mpi4py import MPI
+#from mpi4py import MPI
 
 class Monitor(Element, metaclass=ABCMeta):
     """

tracking/optics.py

@@ -215,6 +215,7 @@ class Optics:
         """
         return self._local_gamma
     
+        
     def beta(self, position):
         """
         Return beta functions at specific locations given by position. If no

tracking/particles.py

@@ -231,6 +231,16 @@ class Bunch:
                  np.mean(self['tau']*self['delta'])**2)**(0.5)
         return np.array([emitX, emitY, emitS])
     
+    @property
+    def cs_invariant(self):
+        Jx = (self.ring.optics.local_gamma * np.mean(self['x']**2)) + \
+             (2*self.ring.optics.local_alpha * np.mean(self['x']*self['xp'])) + \
+             (self.ring.optics.local_beta * np.mean(self['xp']**2))
+        Jy = (self.ring.optics.local_gamma * np.mean(self['y']**2)) + \
+             (2*self.ring.optics.local_alpha * np.mean(self['y']*self['yp'])) + \
+             (self.ring.optics.local_beta * np.mean(self['yp']**2))
+        return np.array((Jx,Jy))
+        
     @property
     def energy_change(self):
         """Store the particle relative energy change in the last turn. Used by

tracking/synchrotron.py

@@ -265,6 +265,21 @@ class Synchrotron:
             sigma[3,:] = (self.emit[1]*self.optics.alpha(position)[1] +
                         self.optics.dispersion(position)[3]**2*self.sigma_delta)**0.5
         return sigma
+    #------------------------------------------
+    # @property
+    # def sigma(self):
+    #     """RMS beam size at equilibrium"""
+    #     sigma = np.zeros((4,))
+    #     sigma[0] = (self.emit[0]*self.optics.beta(10)[0] +
+    #                 self.optics.dispersion(10)[0]**2*self.sigma_delta)**0.5
+    #     sigma[1] = (self.emit[0]*self.optics.alpha(10)[0] +
+    #                 self.optics.dispersion(10)[1]**2*self.sigma_delta)**0.5
+    #     sigma[2] = (self.emit[1]*self.optics.beta(10)[1] +
+    #                 self.optics.dispersion(10)[2]**2*self.sigma_delta)**0.5
+    #     sigma[3] = (self.emit[1]*self.optics.alpha(10)[1] +
+    #                 self.optics.dispersion(10)[3]**2*self.sigma_delta)**0.5
+    #     return sigma
+    #------------------------------------------
     
     def synchrotron_tune(self, Vrf):
         """