Add methods to compute longitudinal Twiss and invariant

Add Synchrotron.get_longitudinal_twiss method to compute the longitudinal Twiss parameters and the synchrotron tune for single or multi-harmonic RF systems.
Change back synchrotron_tune to provide only the usual the (unperturbed) synchrotron tune from main RF voltage.
The longitudinal invariant Js is now computed in Bunch.cs_invariant and Beam.bunch_cs and save by BeamMontior and BunchMonitor, its value is 0 by default if the longitudinal Twiss parameters are not provided in the Synchrotron class.

mbtrack2/tracking/monitors/monitors.py

@@ -289,10 +289,10 @@ class BunchMonitor(Monitor):
         group_name = "BunchData_" + str(self.bunch_number)
         dict_buffer = {"mean":(6, buffer_size), "std":(6, buffer_size),
                      "emit":(3, buffer_size), "current":(buffer_size,),
-                     "cs_invariant":(2, buffer_size)}
+                     "cs_invariant":(3, buffer_size)}
         dict_file = {"mean":(6, total_size), "std":(6, total_size),
                      "emit":(3, total_size), "current":(total_size,),
-                     "cs_invariant":(2, total_size)}
+                     "cs_invariant":(3, total_size)}
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
         
@@ -442,12 +442,12 @@ class BeamMonitor(Monitor):
                        "std" : (6, h, buffer_size),
                        "emit" : (3, h, buffer_size),
                        "current" : (h, buffer_size),
-                       "cs_invariant" : (2, h, buffer_size)}
+                       "cs_invariant" : (3, h, buffer_size)}
         dict_file = {"mean" : (6, h, total_size), 
                        "std" : (6, h, total_size),
                        "emit" : (3, h, total_size),
                        "current" : (h, total_size),
-                       "cs_invariant" : (2, h, total_size)}
+                       "cs_invariant" : (3, h, total_size)}
         
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)

mbtrack2/tracking/monitors/plotting.py

@@ -26,7 +26,7 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
     dimension : str
          The dimension of the dataset to plot:
             for "emit", dimension = {"x","y","s"},
-            for "cs_invariant", dimension = {"x","y"},
+            for "cs_invariant", dimension = {"x","y","s"},
             for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
             not used if "current".
         The default is "tau".
@@ -78,9 +78,9 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
                     y = np.nanstd(data[axis,bunch_index,:],0)
                 y_label = stat_var + " " + label[axis]
             elif dataset == "cs_invariant":
-                dimension_dict = {"x":0, "y":1}
+                dimension_dict = {"x":0, "y":1, "s":2}
                 axis = dimension_dict[dimension]
-                label = ['$J_x$ (m)', '$J_y$ (m)']
+                label = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']
                 if stat_var == "mean":
                     y = np.nanmean(data[axis,bunch_index,:],0)
                 elif stat_var == "std":
@@ -123,9 +123,9 @@ def plot_beamdata(filenames, dataset="mean", dimension="tau", stat_var="mean",
                 y = data[axis,:,idx]
                 y_label = label[axis]
             elif dataset == "cs_invariant":
-                dimension_dict = {"x":0, "y":1}
+                dimension_dict = {"x":0, "y":1, "s":2}
                 axis = dimension_dict[dimension]
-                label = ['$J_x$ (m)', '$J_y$ (m)']
+                label = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']
                 y = data[axis,:,idx]
                 y_label = label[axis]
             elif dataset == "mean" or dataset == "std":
@@ -165,7 +165,7 @@ def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None):
     dimension : str
          The dimension of the dataset to plot:
             for "emit", dimension = {"x","y","s"},
-            for "cs_invariant", dimension = {"x","y"},
+            for "cs_invariant", dimension = {"x","y","s"},
             for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
             not used if "current".
         The default is "tau".
@@ -202,9 +202,9 @@ def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None):
         z_label = label[axis]
         title = z_label
     elif dataset == "cs_invariant":
-        dimension_dict = {"x":0, "y":1}
+        dimension_dict = {"x":0, "y":1, "s":2}
         axis = dimension_dict[dimension]
-        label = ['$J_x$ (m)', '$J_y$ (m)']
+        label = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']
         z = np.array(data["cs_invariant"][axis,:,:]).T
         z_label = label[axis]
         title = z_label
@@ -261,7 +261,7 @@ def plot_bunchdata(filenames, bunch_number, dataset, dimension="x",
         The dimension of the dataset to plot. Use "None" for "current",
         otherwise use the following : 
             for "emit", dimension = {"x","y","s"},
-            for "cs_invariant", dimension = {"x","y"},
+            for "cs_invariant", dimension = {"x","y","s"},
             for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"},
             for "action", dimension = {"x","y"}.
     legend : list of str, optional
@@ -320,10 +320,10 @@ def plot_bunchdata(filenames, bunch_number, dataset, dimension="x",
             label = label_list[axis_index]
             
         elif dataset == "cs_invariant":
-            dimension_dict = {"x":0, "y":1}
+            dimension_dict = {"x":0, "y":1, "s":2}
             axis_index = dimension_dict[dimension]
             y_var = file[group][dataset][axis_index]
-            label_list = ['$J_x$ (m)', '$J_y$ (m)']
+            label_list = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']
             label = label_list[axis_index]
 
         x_axis = file[group]["time"][:]

mbtrack2/tracking/particles.py

@@ -255,7 +255,10 @@ class Bunch:
         Jy = (self.ring.optics.local_gamma[1] * self['y']**2) + \
               (2*self.ring.optics.local_alpha[1] * self['y']*self['yp']) + \
               (self.ring.optics.local_beta[1] * self['yp']**2)
-        return np.array((np.mean(Jx),np.mean(Jy)))
+        Js = (self.ring.long_gamma * self['tau']**2) + \
+              (2*self.ring.long_alpha * self['tau']*self['delta']) + \
+              (self.ring.long_beta * self['delta']**2)
+        return np.array((np.mean(Jx), np.mean(Jy), np.mean(Js)))
         
     def init_gaussian(self, cov=None, mean=None, **kwargs):
         """
@@ -778,7 +781,7 @@ class Beam:
     def bunch_cs(self):
         """Return an array with the average Courant-Snyder invariant for each 
         bunch"""
-        bunch_cs = np.zeros((2,self.ring.h))
+        bunch_cs = np.zeros((3,self.ring.h))
         for idx, bunch in enumerate(self.not_empty):
             index = self.bunch_index[idx]
             bunch_cs[:,index] = bunch.cs_invariant

mbtrack2/tracking/synchrotron.py

@@ -82,11 +82,21 @@ class Synchrotron:
         Relativistic Lorentz gamma.
     beta : float
         Relativistic Lorentz beta.
+    long_alpha : float
+        Longitudinal alpha Twiss parameter at the tracking location.
+        Initialized at zero.
+    long_beta : float
+        Longitudinal beta Twiss parameter at the tracking location in [s].
+        Initialized at zero.
+    long_gamma : float
+        Longitudinal gamma Twiss parameter at the tracking location in [s-1].
+        Initialized at zero.
+        
         
     Methods
     -------
     synchrotron_tune(V)
-        Compute the synchrotron tune for single or multi-harmonic RF systems.
+        Compute the (unperturbed) synchrotron tune from main RF voltage.
     sigma(position)
         Return the RMS beam size at equilibrium in [m].
     eta(delta)
@@ -100,11 +110,17 @@ class Synchrotron:
         componenet from AT lattice.
     get_mcf_order()
         Compute momentum compaction factor up to 3rd order from AT lattice.
+    get_longitudinal_twiss(V)
+        Compute the longitudinal Twiss parameters and the synchrotron tune for 
+        single or multi-harmonic RF systems.
     """
     def __init__(self, h, optics, particle, **kwargs):
         self._h = h
         self.particle = particle
         self.optics = optics
+        self.long_alpha = 0
+        self.long_beta = 0
+        self.long_gamma = 0
         
         if self.optics.use_local_values == False:
             self.L = kwargs.get('L', self.optics.lattice.circumference)
@@ -317,56 +333,22 @@ class Synchrotron:
                         self.optics.dispersion(position)[3]**2*self.sigma_delta**2)**0.5
         return sigma
     
-    def synchrotron_tune(self, V, phase=None, harmonics=None):
+    def synchrotron_tune(self, V):
         """
-        Compute the synchrotron tune for single or multi-harmonic RF systems.
-        
-        Hypthesis:        
-        - Use the linear approximation (tau ~ 0).
-        - For higher haromics only, cos(phi) ~ 0 must be true.
+        Compute the (unperturbed) synchrotron tune from main RF voltage.
         
         Parameters
         ----------
-        V : float or array-like of float
-            RF voltage in [V].
-        phase : float or array-like of float, optional
-            RF phase using cosine convention in [rad].
-            Default is None.
-        harmonics : int or array-like of int, optional
-            Harmonics of the cavities to consider.
-            Default is None.
-            
-        Usage
-        -----
-        - If a single voltage value is given, compute the synchrotron tune for 
-        a single RF system set at the synchronous phase.
-        - If a single pair of voltage/phase value is given, compute the 
-        synchrotron tune for a single RF system.
-        - Otherwise, compute the synchrotron tune for a multi-harmonic RF 
-        system.
-            
+        V : float
+            Main RF voltage in [V].
+
         Returns
         -------
         tuneS : float
-            Synchrotron tune in the linear approximation.
+            Synchrotron tune.
+            
         """
-        
-        if isinstance(V, float) or isinstance(V, int):
-            V = [V]
-            if phase is None:
-                phase = [np.arccos(self.U0/V[0])]
-            elif isinstance(phase, float) or isinstance(phase, int):
-                phase = [phase]
-            if harmonics is None:
-                harmonics = [1]
-                
-        if not (len(V) == len(phase) == len(harmonics)):
-            raise ValueError("You must provide array of the same length for"
-                             " V, phase and harmonics")
-        
-        Vsum = 0
-        for i in range(len(V)):
-            Vsum += harmonics[i] * V[i] * np.sin(phase[i])
+        Vsum =  V * np.sin(np.arccos(self.U0/V))
         phi = np.arccos(1 - self.eta(0) * np.pi * self.h / self.E0 * Vsum)
         tuneS = phi / (2 * np.pi)
         return tuneS
@@ -438,3 +420,75 @@ class Synchrotron:
             self.mcf_order = pvalue
         else:
             return pvalue
+
+    def get_longitudinal_twiss(self, V, phase=None, harmonics=None, add=True):
+        """
+        Compute the longitudinal Twiss parameters and the synchrotron tune for 
+        single or multi-harmonic RF systems.
+        
+        Hypthesis:
+        - Use the linear approximation (tau ~ 0).
+        - For higher haromics only, cos(phi) ~ 0 must be true.
+        
+        Parameters
+        ----------
+        V : float or array-like of float
+            RF voltage in [V].
+        phase : float or array-like of float, optional
+            RF phase using cosine convention in [rad].
+            Default is None.
+        harmonics : int or array-like of int, optional
+            Harmonics of the cavities to consider.
+            Default is None.
+        add : bool, optional
+            If True, 
+            
+        Usage
+        -----
+        - If a single voltage value is given, assume a single RF system set at 
+        the synchronous phase.
+        - If a single pair of voltage/phase value is given, assume a single RF 
+        system at this setting.
+        - Otherwise, compute the synchrotron tune for a multi-harmonic RF 
+        system.
+            
+        Returns
+        -------
+        tuneS : float
+            Synchrotron tune in the linear approximation.
+        long_alpha : float
+            Longitudinal alpha Twiss parameter at the tracking location.
+        long_beta : float
+            Longitudinal beta Twiss parameter at the tracking location in [s].   
+        long_gamma : float
+            Longitudinal gamma Twiss parameter at the tracking location in [s-1].
+        """
+        
+        if isinstance(V, float) or isinstance(V, int):
+            V = [V]
+            if phase is None:
+                phase = [np.arccos(self.U0/V[0])]
+            elif isinstance(phase, float) or isinstance(phase, int):
+                phase = [phase]
+            if harmonics is None:
+                harmonics = [1]
+                
+        if not (len(V) == len(phase) == len(harmonics)):
+            raise ValueError("You must provide array of the same length for"
+                             " V, phase and harmonics")
+        
+        Vsum = 0
+        for i in range(len(V)):
+            Vsum += harmonics[i] * V[i] * np.sin(phase[i])
+        phi = np.arccos(1 - self.eta(0) * np.pi * self.h / self.E0 * Vsum)
+        long_alpha = - self.eta(0) * np.pi * self.h / (self.E0 * np.sin(phi)) * Vsum
+        long_beta = self.eta(0) * self.T0 / np.sin(phi)
+        long_gamma = self.omega1 * Vsum / (self.E0 * np.sin(phi))
+        tuneS = phi / (2 * np.pi)
+        if add:
+            self.tuneS = tuneS
+            self.long_alpha = long_alpha
+            self.long_beta = long_beta
+            self.long_gamma = long_gamma
+        else:
+            return tuneS, long_alpha, long_beta, long_gamma
\ No newline at end of file