diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py index 7e36e819e0b56b3720e6bd444a021405692242d7..049f7006f60c3f844e7c50c391204b9816a0c8ca 100644 --- a/mbtrack2/tracking/element.py +++ b/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) diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py index c723eea002d3c5bff94bdd52aeb2075e609dc930..b41093411771ad3f5ef9004141183477705e3534 100644 --- a/mbtrack2/tracking/synchrotron.py +++ b/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): """