Include ADTS in phase advance calculation

tracking/element.py

@@ -153,6 +153,11 @@ class TransverseMap(Element):
         self.gamma = self.ring.optics.local_gamma
         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:
+            self.adts_option = True
     
     @Element.parallel    
     def track(self, bunch):
@@ -166,9 +171,17 @@ class TransverseMap(Element):
         bunch : Bunch or Beam object
         """
 
-        # Compute phase adcence which depends on energy via chromaticity
-        phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])
-        phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
+        # 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']))
+            
+            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']))
+        
+        else:
+            phase_advance_x = self.phase_advance[0]*(1+self.ring.chro[0]*bunch["delta"])
+            phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
         
         # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
         matrix = np.zeros((6,6,len(bunch)))

tracking/synchrotron.py

@@ -99,6 +99,7 @@ class Synchrotron:
         self.sigma_delta = kwargs.get('sigma_delta') # Equilibrium energy spread
         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)
                 
     @property
     def h(self):