Merge branch 'TransverseMap' into develop

mbtrack2/tracking/element.py

@@ -6,7 +6,6 @@ included in the tracking.
 """
 
 import numpy as np
-import at
 from abc import ABCMeta, abstractmethod
 from functools import wraps
 from copy import deepcopy
@@ -176,14 +175,20 @@ class TransverseMap(Element):
             phase_advance_y = 2*np.pi * (self.ring.tune[1] + 
                                          self.ring.chro[1]*bunch["delta"])
         else:
+            Jx = (self.ring.optics.local_gamma[0] * bunch['x']**2) + \
+                  (2*self.ring.optics.local_alpha[0] * bunch['x']*bunch['xp']) + \
+                  (self.ring.optics.local_beta[0] * bunch['xp']**2)
+            Jy = (self.ring.optics.local_gamma[1] * bunch['y']**2) + \
+                  (2*self.ring.optics.local_alpha[1] * bunch['y']*bunch['yp']) + \
+                  (self.ring.optics.local_beta[1] * bunch['yp']**2)
             phase_advance_x = 2*np.pi * (self.ring.tune[0] + 
                                          self.ring.chro[0]*bunch["delta"] + 
-                                         self.adts_poly[0](bunch['x']) + 
-                                         self.adts_poly[2](bunch['y']))
+                                         self.adts_poly[0](Jx) + 
+                                         self.adts_poly[2](Jy))
             phase_advance_y = 2*np.pi * (self.ring.tune[1] + 
                                          self.ring.chro[1]*bunch["delta"] +
-                                         self.adts_poly[1](bunch['x']) + 
-                                         self.adts_poly[3](bunch['y']))
+                                         self.adts_poly[1](Jx) + 
+                                         self.adts_poly[3](Jy))
         
         # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
         matrix = np.zeros((6,6,len(bunch)))
@@ -282,9 +287,11 @@ class TransverseMapSector(Element):
         self.ring = ring
         self.alpha0 = alpha0
         self.beta0 = beta0
+        self.gamma0 = (1 + self.alpha0**2)/self.beta0
         self.dispersion0 = dispersion0
         self.alpha1 = alpha1
         self.beta1 = beta1
+        self.gamma1 = (1 + self.alpha1**2)/self.beta1
         self.dispersion1 = dispersion1  
         self.tune_diff = phase_diff / (2*np.pi)
         self.chro_diff = chro_diff
@@ -315,14 +322,20 @@ class TransverseMapSector(Element):
             phase_advance_y = 2*np.pi * (self.tune_diff[1] + 
                                          self.chro_diff[1]*bunch["delta"])
         else:
+            Jx = (self.gamma0[0] * bunch['x']**2) + \
+                  (2*self.alpha0[0] * bunch['x']*self['xp']) + \
+                  (self.beta0[0] * bunch['xp']**2)
+            Jy = (self.gamma0[1] * bunch['y']**2) + \
+                  (2*self.alpha0[1] * bunch['y']*bunch['yp']) + \
+                  (self.beta0[1] * bunch['yp']**2)
             phase_advance_x = 2*np.pi * (self.tune_diff[0] + 
                                          self.chro_diff[0]*bunch["delta"] + 
-                                         self.adts_poly[0](bunch['x']) + 
-                                         self.adts_poly[2](bunch['y']))
+                                         self.adts_poly[0](Jx) + 
+                                         self.adts_poly[2](Jy))
             phase_advance_y = 2*np.pi * (self.tune_diff[1] + 
                                          self.chro_diff[1]*bunch["delta"] +
-                                         self.adts_poly[1](bunch['x']) + 
-                                         self.adts_poly[3](bunch['y']))
+                                         self.adts_poly[1](Jx) + 
+                                         self.adts_poly[3](Jy))
         
         # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
         matrix = np.zeros((6,6,len(bunch)))
@@ -379,7 +392,7 @@ def transverse_map_sector_generator(ring, positions):
         List of TransverseMapSector elements.
 
     """
-    
+    import at
     def _compute_chro(ring, pos, dp=1e-4):
         lat = deepcopy(ring.optics.lattice)
         lat.append(at.Marker("END"))

mbtrack2/tracking/synchrotron.py

@@ -47,13 +47,13 @@ class Synchrotron:
         The order of the elements strictly needs to be
         [coef_xx, coef_yx, coef_xy, coef_yy], where x and y denote the horizontal
         and the vertical plane, respectively, and coef_PQ means the polynomial's
-        coefficients of the ADTS in plane P due to the offset in plane Q.
+        coefficients of the ADTS in plane P due to the amplitude in plane Q.
 
-        For example, if the tune shift in y due to the offset x is characterized
-        by the equation dQy(x) = 3*x**2 + 2*x + 1, then coef_yx takes the form
-        np.array([3, 2, 1]).
+        For example, if the tune shift in y due to the Courant-Snyder invariant 
+        Jx is characterized by the equation dQy(x) = 3*Jx**2 + 2*Jx + 1, then 
+        coef_yx takes the form np.array([3, 2, 1]).
         
-        Use None, to exclude the ADTS calculation.
+        Use None, to exclude the ADTS from calculation.
         
     Attributes
     ----------
@@ -296,3 +296,20 @@ class Synchrotron:
         tuneS = np.sqrt( - (Vrf / self.E0) * (self.h * self.ac) / (2*np.pi) 
                         * np.cos(phase) )
         return tuneS
+    
+    def get_adts(self):
+        """
+        Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar 
+        componenet from AT lattice.
+        """
+        import at
+        if self.optics.use_local_values:
+            raise ValueError("ADTS needs to be provided manualy as no AT" + 
+                             " lattice file is loaded.")
+            
+        det = at.physics.nonlinear.gen_detuning_elem(self.optics.lattice)
+        coef_xx = np.array([det.A1/2, 0])
+        coef_yx = np.array([det.A2/2, 0])
+        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]