changed chromaticity computation in _compute_chro() of...

changed chromaticity computation in _compute_chro() of transverse_map_sector_generator(). Now it can compute arbitrary order of chromaticity from AT.

mbtrack2/tracking/element.py

@@ -404,22 +404,15 @@ def transverse_map_sector_generator(ring, positions):
     else:
         import at
 
-        def _compute_chro(ring, pos, dp=1e-4):
+        def _compute_chro(ring, pos, dp=1e-4, order=4):
             lat = deepcopy(ring.optics.lattice)
             lat.append(at.Marker("END"))
-            N = len(lat)
-            refpts = np.arange(N)
-            (*elem_neg_dp, ) = at.linopt2(lat, refpts=refpts, dp=-dp)
-            (*elem_pos_dp, ) = at.linopt2(lat, refpts=refpts, dp=dp)
-
-            s = elem_neg_dp[2]["s_pos"]
-            mux0 = elem_neg_dp[2]["mu"][:, 0]
-            mux1 = elem_pos_dp[2]["mu"][:, 0]
-            muy0 = elem_neg_dp[2]["mu"][:, 1]
-            muy1 = elem_pos_dp[2]["mu"][:, 1]
-
-            Chrox = (mux1-mux0) / (2*dp) / 2 / np.pi
-            Chroy = (muy1-muy0) / (2*dp) / 2 / np.pi
+            fit, dpa, tune = at.physics.nonlinear.chromaticity(lat,
+                                                               method='linopt',
+                                                               dpm=0.02,
+                                                               n_points=100,
+                                                               order=order)
+            Chrox, Chroy = fit
             chrox = np.interp(pos, s, Chrox)
             chroy = np.interp(pos, s, Chroy)