diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index 0f69f6cfb391776d5f747b2230fb9a9c374df895..1575a1973bc4a1ddedc33b0318f4f464f95ff051 100644
--- a/mbtrack2/tracking/element.py
+++ b/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])