Merge branch 'fix-transverse_map_sector_generator' into 'develop'

Fix transverse map sector generator

See merge request !38

mbtrack2/tracking/element.py

@@ -415,7 +415,7 @@ class TransverseMap(TransverseMapSector):
                          2 * np.pi * ring.tune, ring.chro, ring.adts)
 
 
-def transverse_map_sector_generator(ring, positions):
+def transverse_map_sector_generator(ring, positions, **kwargs):
     """
     Convenience function which generate a list of TransverseMapSector elements
     from a ring:
@@ -435,6 +435,8 @@ def transverse_map_sector_generator(ring, positions):
         of the TransverseMapSector elements.
         The array should contain the initial position (s=0) but not the end
         position (s=ring.L), so like position = np.array([0, pos1, pos2, ...]).
+    
+    See at.physics.nonlinear.chromaticity for **kwargs
 
     Returns
     -------
@@ -463,8 +465,10 @@ def transverse_map_sector_generator(ring, positions):
                                     adts=adts))
     else:
         import at
+        dp = kwargs.get('dp', 1e-2)
+        order = kwargs.get('order', 1)
 
-        def _compute_chro(ring, N_sec, dp=1e-2, order=4):
+        def _compute_chro(ring, N_sec, dp, order):
             lat = deepcopy(ring.optics.lattice)
             lat.append(at.Marker("END"))
             fit, _, _ = at.physics.nonlinear.chromaticity(lat,
@@ -472,12 +476,18 @@ def transverse_map_sector_generator(ring, positions):
                                                           dpm=dp,
                                                           n_points=100,
                                                           order=order)
-            Chrox, Chroy = fit[0, 1:], fit[1, 1:]
-            chrox = np.array([np.linspace(0, val, N_sec) for val in Chrox])
-            chroy = np.array([np.linspace(0, val, N_sec) for val in Chroy])
-            return np.array([chrox, chroy])
+            chro_xy = [
+                elem for pair in zip(fit[0, 1:], fit[1, 1:]) for elem in pair
+            ]
+            len_chro = int(order * 2)
+            _chro = np.zeros((len_chro, N_sec))
+            for i in range(len_chro):
+                chro_order_splited = np.linspace(0, chro_xy[i], N_sec)
+                _chro[i, :] = chro_order_splited
+
+            return _chro
 
-        _chro = _compute_chro(ring, N_sec)
+        _chro = _compute_chro(ring, N_sec, dp, order)
         for i in range(N_sec):
             alpha0 = ring.optics.alpha(positions[i])
             beta0 = ring.optics.beta(positions[i])