Merge branch 'export_pyat' into develop

mbtrack2/tracking/synchrotron.py

@@ -114,6 +114,8 @@ class Synchrotron:
     get_longitudinal_twiss(V)
         Compute the longitudinal Twiss parameters and the synchrotron tune for 
         single or multi-harmonic RF systems.
+    to_pyat(Vrf)
+        Return a pyAT simple_ring element from the Synchrotron element data.
     """
     def __init__(self, h, optics, particle, **kwargs):
         self._h = h
@@ -506,3 +508,77 @@ class Synchrotron:
             self.long_gamma = long_gamma
         else:
             return tuneS, long_alpha, long_beta, long_gamma
+        
+    def to_pyat(self, Vrf, harmonic_number=None, TimeLag=False):
+        """
+        Return a pyAT simple_ring element from the Synchrotron element data.
+        
+        See pyAT documentation for informations about simple_ring.
+
+        Parameters
+        ----------
+        Vrf : float or array-like of float
+            RF Voltage in [V]. If sevral cavities are provided, harmonic_number
+            and TimeLag should be provided.
+        harmonic_number : float or array-like of float, optional
+            Harmonic numbers of the RF cavities. The default is None.
+        TimeLag : float or array-like of float, optional
+            Set the timelag of the cavities in pyAT definition. 
+            The default is False.
+
+        Returns
+        -------
+        at_simple_ring : at.physics.fastring.simple_ring
+            A pyAT simple_ring element.
+
+        """
+        from at import simple_ring
+        optics = self.optics
+        
+        if (harmonic_number is None) and isinstance(Vrf, (float, int)):
+            harmonic_number=self.h
+        
+        if isinstance(Vrf, (list, np.ndarray)):
+            if (harmonic_number is None) or (TimeLag is None):
+                raise ValueError("If sevral cavities are provided, "
+                         "harmonic_number and TimeLag should be provided.")
+                
+        if self.adts is not None:
+            try:
+                A1 = self.adts[0][-2]*2
+                A2 = self.adts[1][-2]*2
+                A3 = self.adts[3][-2]*2
+            except IndexError:
+                A1 = None
+                A2 = None
+                A3 = None
+        else: 
+            A1 = None
+            A2 = None
+            A3 = None
+        
+        at_simple_ring = simple_ring(energy=self.E0,
+                                     circumference=self.L,
+                                     harmonic_number=harmonic_number,
+                                     Qx=self.tune[0],
+                                     Qy=self.tune[1],
+                                     Vrf=Vrf,
+                                     alpha=self.ac,
+                                     betax=optics.local_beta[0],
+                                     betay=optics.local_beta[1],
+                                     alphax=optics.local_alpha[0],
+                                     alphay=optics.local_alpha[1],
+                                     Qpx=self.chro[0],
+                                     Qpy=self.chro[1],
+                                     A1=A1,
+                                     A2=A2,
+                                     A3=A3,
+                                     emitx=self.emit[0],
+                                     emity=self.emit[1],
+                                     espread=self.sigma_delta,
+                                     taux=self.tau[0]/self.T0,
+                                     tauy=self.tau[1]/self.T0,
+                                     tauz=self.tau[2]/self.T0,
+                                     U0=self.U0,
+                                     TimeLag=TimeLag)
+        return at_simple_ring