bugfixes

mbtrack2/tracking/element.py

@@ -404,15 +404,15 @@ def transverse_map_sector_generator(ring, positions):
     else:
         import at
 
-        def _compute_chro(ring, pos, dp=1e-4, order=4):
+        def _compute_chro(ring, pos, dp=1e-2, order=4):
             lat = deepcopy(ring.optics.lattice)
             lat.append(at.Marker("END"))
             fit, dpa, tune = at.physics.nonlinear.chromaticity(lat,
                                                                method='linopt',
-                                                               dpm=0.02,
+                                                               dpm=dp,
                                                                n_points=100,
                                                                order=order)
-            Chrox, Chroy = fit
+            Chrox, Chroy = fit[0, 1:], fit[1, 1:]
             chrox = np.interp(pos, s, Chrox)
             chroy = np.interp(pos, s, Chroy)
 

mbtrack2/tracking/synchrotron.py

@@ -388,15 +388,20 @@ class Synchrotron:
         self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]
 
     def get_chroma(self, order=4, dpm=0.02, n_points=100):
+        """
+        Compute and add chromaticity (linear and nonlinear) from AT lattice and update the property.
+        """
         import at
-        fit, dpa, tune = at.physics.nonlinear.chromaticity(
-            self.ring.optics.lattice,
-            method='linopt',
-            dpm=0.02,
-            n_points=100,
-            order=order)
+        fit, dpa, tune = at.physics.nonlinear.chromaticity(self.optics.lattice,
+                                                           method='linopt',
+                                                           dpm=dpm,
+                                                           n_points=n_points,
+                                                           order=order)
         chrox, chroy = fit
-        return chrox, chroy
+        self.chro = [
+            elem for pair in zip(chrox[1:], chroy[1:]) for elem in pair
+        ]
+        return chrox[1:], chroy[1:]
 
     def get_mcf_order(self, add=True, show_fit=False):
         """