feature-improve Synchrotron.get_adts

Adds more control on the underneath pyAT function and optional plotting to check result.

mbtrack2/tracking/synchrotron.py

@@ -373,23 +373,78 @@ class Synchrotron:
         tuneS = phi / (2 * np.pi)
         return tuneS
 
-    def get_adts(self):
+    def get_adts(self,
+                 xm=1e-4,
+                 ym=1e-4,
+                 npoints=9,
+                 plot=False,
+                 ax=None,
+                 **kwargs):
         """
         Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar
         componenet from AT lattice.
+        
+        Parameters
+        ----------
+        xm : float, optional
+            Maximum horizontal amplitude in [m].
+            Default is 1e-4.
+        ym : float, optional
+            Maximum vertical amplitude in [m].
+            Default is 1e-4.
+        npoints : int, optional
+            Number of points in each plane.
+            Default is 9.
+        plot : bool, optional
+            If True, plot the computed tune shift with amplitude.
+            Default is False.
+        ax : array of shape (2,) of matplotlib.axes.Axes, optional
+            Axes where to plot the figures.
+            Default is None.
+            
+        Returns
+        -------
+        ax : array of shape (2,) of matplotlib.axes.Axes
+            Only if plot is True.
+        
+        See at.physics.nonlinear.detuning for **kwargs
         """
         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])
+        r0, r1, x, q_dx, y, q_dy = at.physics.nonlinear.detuning(
+            self.optics.lattice.radiation_off(copy=True),
+            npoints=npoints,
+            xm=xm,
+            ym=ym,
+            **kwargs)
+        coef_xx = np.array([r1[0][0] / 2, 0])
+        coef_yx = np.array([r1[0][1] / 2, 0])
+        coef_xy = np.array([r1[0][1] / 2, 0])
+        coef_yy = np.array([r1[1][1] / 2, 0])
         self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]
 
+        if plot:
+            if ax is None:
+                fig, ax = plt.subplots(2, 1)
+
+            ax[0].scatter(x * 1e6, q_dx[:, 0])
+            ax[0].scatter(y * 1e6, q_dy[:, 0])
+            ax[0].set_ylabel("Delta Tune X")
+            ax[0].set_xlabel("Amplitude [um]")
+            ax[0].legend(['dqx/dx', 'dqx/dy'])
+            ax[0].grid()
+
+            ax[1].scatter(x * 1e6, q_dx[:, 1])
+            ax[1].scatter(y * 1e6, q_dy[:, 1])
+            ax[1].set_ylabel("Delta Tune Y")
+            ax[1].set_xlabel("Amplitude [um]")
+            ax[1].legend(['dqy/dx', 'dqy/dy'])
+            ax[1].grid()
+
+            return ax
+
     def get_chroma(self, order=4, dpm=0.02, n_points=100):
         """
         Compute chromaticity (linear and nonlinear) from AT lattice and update the property.

tests/unit/tracking/test_synchrotron.py

@@ -109,6 +109,10 @@ class TestSynchrotron:
         ring_with_at_lattice.get_adts()
         assert ring_with_at_lattice.adts is not None
         
+    def test_get_adts_plot(self, ring_with_at_lattice):
+        axes = ring_with_at_lattice.get_adts(plot=True)
+        assert axes is not None
+        
     def test_get_chroma(self, ring_with_at_lattice):
         ring_with_at_lattice.get_chroma()
         assert len(ring_with_at_lattice.chro) == 8