Fix ADTS handling algorithm and add its docstring

tracking/element.py

@@ -154,10 +154,15 @@ class TransverseMap(Element):
         self.dispersion = self.ring.optics.local_dispersion
         self.phase_advance = self.ring.tune[0:2]*2*np.pi
         
-        if self.ring.adts.all() == 0:
-            self.adts_option = False
-        else:
+        if self.ring.adts is not None:
             self.adts_option = True
+            self.adts_poly = np.array([np.poly1d(self.ring.adts[0]),
+                                       np.poly1d(self.ring.adts[1]),
+                                       np.poly1d(self.ring.adts[2]), 
+                                       np.poly1d(self.ring.adts[3])])
+        
+        else:
+            self.adts_option = False
     
     @Element.parallel    
     def track(self, bunch):
@@ -174,10 +179,10 @@ class TransverseMap(Element):
         # Compute phase advance which depends on energy via chromaticity and ADTS
         if self.adts_option is True:
             phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"]+
-                                self.ring.adts[0](bunch['x'])+self.ring.adts[2](bunch['y']))
+                              self.adts_poly[0](bunch['x'])+self.adts_poly[2](bunch['y']))
             
             phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"]+
-                                self.ring.adts[1](bunch['x'])+self.ring.adts[3](bunch['y']))
+                              self.adts_poly[1](bunch['x'])+self.adts_poly[3](bunch['y']))
         
         else:
             phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])

tracking/synchrotron.py

@@ -44,6 +44,19 @@ class Synchrotron:
         Horizontal and vertical (non-normalized) chromaticities.
     U0 : float, optional
         Energy loss per turn in [eV].
+    adts : array of shape (4,) or None, optional
+        Array that contains arrays of polynomial's coefficients, in decreasing 
+        powers, used to determine Amplitude-Dependent Tune Shifts (ADTS). 
+        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.
+
+        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]).
+        
+        Use None, to exclude the ADTS calculation.
         
     Attributes
     ----------