Add higher order momentum compaction factor

Add mcf_order optional parameter for Synchrotron class to prodivde higher-orders momentum compaction factor.
Add Synchrotron.get_mcf_order method to compute momentum compaction factor up to 3rd order from AT lattice.
Modify Synchrotron.eta to use higher order mcf.
Improve Synchrotron docstrings.
Change LongitudinalMap to use Synchrotron.eta.

mbtrack2/tracking/element.py

@@ -91,7 +91,7 @@ class LongitudinalMap(Element):
         bunch : Bunch or Beam object
         """
         bunch["delta"] -= self.ring.U0 / self.ring.E0
-        bunch["tau"] += self.ring.ac * self.ring.T0 * bunch["delta"]
+        bunch["tau"] += self.ring.eta(bunch["delta"]) * self.ring.T0 * bunch["delta"]
 
 class SynchrotronRadiation(Element):
     """

mbtrack2/tracking/synchrotron.py

@@ -4,6 +4,7 @@ Module where the Synchrotron class is defined.
 """
 
 import numpy as np
+import matplotlib.pyplot as plt
 from scipy.constants import c, e
         
 class Synchrotron:
@@ -54,6 +55,12 @@ class Synchrotron:
         coef_yx takes the form np.array([3, 2, 1]).
         
         Use None, to exclude the ADTS from calculation.
+    mcf_order : array, optional
+        Higher-orders momentum compaction factor in decreasing powers.
+        Input here the coefficients considering only the derivatives and which
+        do not include any factorial factor.
+        The 1st order should be included in the array and be at the last 
+        position.
         
     Attributes
     ----------
@@ -75,16 +82,24 @@ class Synchrotron:
         Relativistic Lorentz gamma.
     beta : float
         Relativistic Lorentz beta.
-    eta : float
-        Momentum compaction.
-    sigma : array of shape (4,)
-        RMS beam size at equilibrium in [m].
         
     Methods
     -------
     synchrotron_tune(Vrf)
         Compute synchrotron tune from RF voltage.
     sigma(position)
+        Return the RMS beam size at equilibrium in [m].
+    eta(delta)
+        Momentum compaction taking into account higher orders if provided in
+        mcf_order.
+    mcf(delta)
+        Momentum compaction factor taking into account higher orders if 
+        provided in mcf_order.
+    get_adts()
+        Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar 
+        componenet from AT lattice.
+    get_mcf_order()
+        Compute momentum compaction factor up to 3rd order from AT lattice.
     """
     def __init__(self, h, optics, particle, **kwargs):
         self._h = h
@@ -111,7 +126,8 @@ class Synchrotron:
         self.sigma_0 = kwargs.get('sigma_0') # Natural bunch length [s]
         self.emit = kwargs.get('emit') # X/Y emittances in [m.rad]
         self.adts = kwargs.get('adts') # Amplitude-Dependent Tune Shift (ADTS)
-                
+        self.mcf_order = kwargs.get('mcf_order', self.ac) # Higher-orders momentum compaction factor
+        
     @property
     def h(self):
         """Harmonic number"""
@@ -231,9 +247,32 @@ class Synchrotron:
         self.gamma = value/(self.particle.mass*c**2/e)
 
     @property
-    def eta(self):
-        """Momentum compaction"""
-        return self.ac - 1/(self.gamma**2)
+    def mcf_order(self):
+        """Higher-orders momentum compaction factor"""
+        return self._mcf_order
+    
+    @mcf_order.setter
+    def mcf_order(self, value):
+        self._mcf_order = value
+        self.mcf = np.poly1d(self.mcf_order)
+        
+    def eta(self, delta=0):
+        """
+        Momentum compaction taking into account higher orders if provided in
+        mcf_order.
+
+        Parameters
+        ----------
+        delta : float or array like bunch["delta"], optional
+            Relative energy deviation. The default is 0.
+
+        Returns
+        -------
+        float or array
+            Momentum compaction.
+
+        """
+        return self.mcf(delta) - 1/(self.gamma**2)
     
     def sigma(self, position=None):
         """
@@ -313,3 +352,54 @@ class Synchrotron:
         coef_xy = np.array([det.A2/2, 0])
         coef_yy = np.array([det.A3/2, 0])
         self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]
+        
+    def get_mcf_order(self, add=True, show_fit=False):
+        """
+        Compute momentum compaction factor up to 3rd order from AT lattice.
+        
+        Parameters
+        ----------
+        add : bool, optional
+            If True, add the computed coefficient to the Synchrotron object.
+            If False, the result is returned.
+            The default is True.
+        show_fit : bool, optional
+            If True, show a plot with the data and fit.
+            The default is False.
+            
+        Return
+        ------
+        pvalue : array, optional
+            Momentum compaction factor up to 3rd order
+        """
+        import at
+        if self.optics.use_local_values:
+            raise ValueError("ADTS needs to be provided manualy as no AT" + 
+                             " lattice file is loaded.")
+        deltamin = -1e-4
+        deltamax = 1e-4
+
+        delta = np.linspace(deltamin, deltamax, 7)
+        delta4eval = np.linspace(deltamin, deltamax, 21)
+        alpha = np.zeros_like(delta)
+
+        for i in range(len(delta)):
+            alpha[i] = at.physics.revolution.get_mcf(self.optics.lattice, 
+                                                     delta[i])
+        pvalue = np.polyfit(delta, alpha, 2)
+
+        if show_fit:
+            pvalue = np.polyfit(delta, alpha, 2)
+            palpha = np.polyval(pvalue, delta4eval)
+    
+            plt.plot(delta*100, alpha,'k.')
+            plt.grid()
+            plt.plot(delta4eval*100,palpha,'r')
+            plt.xlabel('Energy (%)')
+            plt.ylabel('Momemtum compaction factor')
+            plt.legend(['Data', 'polyfit'])
+            
+        if add:
+            self.mcf_order = pvalue
+        else:
+            return pvalue