Split correctly the chromaticity

Now split correctly the chromaticity in transverse_map_sector_generator function.
Clarify usage for transverse_map_sector_generator.

mbtrack2/tracking/element.py

@@ -6,8 +6,10 @@ included in the tracking.
 """
 
 import numpy as np
+import at
 from abc import ABCMeta, abstractmethod
 from functools import wraps
+from copy import deepcopy
 from mbtrack2.tracking.particles import Beam
 
 class Element(metaclass=ABCMeta):
@@ -368,6 +370,8 @@ def transverse_map_sector_generator(ring, positions):
     positions : array
         List of longitudinal positions in [m] to use as starting and end points
         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, ...]).
 
     Returns
     -------
@@ -376,28 +380,52 @@ def transverse_map_sector_generator(ring, positions):
 
     """
     
+    def _compute_chro(ring, pos, dp=1e-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
+        chrox = np.interp(pos, s, Chrox)
+        chroy = np.interp(pos, s, Chroy)
+        
+        return np.array([chrox, chroy])
+    
     if ring.optics.use_local_values:
         raise ValueError("The Synchrotron object must be loaded from an AT lattice")
     
     N_sec = len(positions)
     sectors = []
-    chro_diff = ring.chro/N_sec
     for i in range(N_sec):
         alpha0 = ring.optics.alpha(positions[i])
         beta0 = ring.optics.beta(positions[i])
         dispersion0 = ring.optics.dispersion(positions[i])
         mu0 = ring.optics.mu(positions[i])
+        chro0 = _compute_chro(ring, positions[i])
         if i != (N_sec - 1):
             alpha1 = ring.optics.alpha(positions[i+1])
             beta1 = ring.optics.beta(positions[i+1])
             dispersion1 = ring.optics.dispersion(positions[i+1])
             mu1 = ring.optics.mu(positions[i+1])
+            chro1 = _compute_chro(ring, positions[i+1])
         else:
             alpha1 = ring.optics.alpha(positions[0])
             beta1 = ring.optics.beta(positions[0])
             dispersion1 = ring.optics.dispersion(positions[0])
             mu1 = ring.optics.mu(ring.L)
+            chro1 = _compute_chro(ring, ring.L)
         phase_diff = mu1 - mu0
+        chro_diff = chro1 - chro0
         sectors.append(TransverseMapSector(ring, alpha0, beta0, dispersion0, 
                      alpha1, beta1, dispersion1, phase_diff, chro_diff))
     return sectors