diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index 4385b57d2f44e87d82648d2ba97831c22051f581..049f7006f60c3f844e7c50c391204b9816a0c8ca 100644
--- a/mbtrack2/tracking/element.py
+++ b/mbtrack2/tracking/element.py
@@ -10,6 +10,7 @@ from copy import deepcopy
 from functools import wraps
 
 import numpy as np
+from scipy.special import factorial
 
 from mbtrack2.tracking.particles import Beam
 
@@ -128,128 +129,26 @@ class SynchrotronRadiation(Element):
         ----------
         bunch : Bunch or Beam object
         """
+        N = len(bunch)
         if self.switch[0] == True:
-            rand = np.random.normal(size=len(bunch))
+            rand = np.random.standard_normal(size=N)
             bunch["delta"] = (1 - 2 * self.ring.T0 / self.ring.tau[2]) * bunch[
                 "delta"] + 2 * self.ring.sigma_delta * (
                     self.ring.T0 / self.ring.tau[2])**0.5 * rand
 
         if self.switch[1] == True:
-            rand = np.random.normal(size=len(bunch))
+            rand = np.random.standard_normal(size=N)
             bunch["xp"] = (1 - 2 * self.ring.T0 / self.ring.tau[0]
                            ) * bunch["xp"] + 2 * self.ring.sigma()[1] * (
                                self.ring.T0 / self.ring.tau[0])**0.5 * rand
 
         if self.switch[2] == True:
-            rand = np.random.normal(size=len(bunch))
+            rand = np.random.standard_normal(size=N)
             bunch["yp"] = (1 - 2 * self.ring.T0 / self.ring.tau[1]
                            ) * bunch["yp"] + 2 * self.ring.sigma()[3] * (
                                self.ring.T0 / self.ring.tau[1])**0.5 * rand
 
 
-class TransverseMap(Element):
-    """
-    Transverse map for a single turn in the synchrotron.
-
-    Parameters
-    ----------
-    ring : Synchrotron object
-    """
-
-    def __init__(self, ring):
-        self.ring = ring
-        self.alpha = self.ring.optics.local_alpha
-        self.beta = self.ring.optics.local_beta
-        self.gamma = self.ring.optics.local_gamma
-        self.dispersion = self.ring.optics.local_dispersion
-        if self.ring.adts is not None:
-            self.adts_poly = [
-                np.poly1d(self.ring.adts[0]),
-                np.poly1d(self.ring.adts[1]),
-                np.poly1d(self.ring.adts[2]),
-                np.poly1d(self.ring.adts[3]),
-            ]
-
-    @Element.parallel
-    def track(self, bunch):
-        """
-        Tracking method for the element.
-        No bunch to bunch interaction, so written for Bunch objects and
-        @Element.parallel is used to handle Beam objects.
-
-        Parameters
-        ----------
-        bunch : Bunch or Beam object
-        """
-
-        # Compute phase advance which depends on energy via chromaticity and ADTS
-        if self.ring.adts is None:
-            phase_advance_x = (
-                2 * np.pi *
-                (self.ring.tune[0] + self.ring.chro[0] * bunch["delta"]))
-            phase_advance_y = (
-                2 * np.pi *
-                (self.ring.tune[1] + self.ring.chro[1] * bunch["delta"]))
-        else:
-            Jx = ((self.ring.optics.local_gamma[0] * bunch["x"]**2) +
-                  (2 * self.ring.optics.local_alpha[0] * bunch["x"] *
-                   bunch["xp"]) +
-                  (self.ring.optics.local_beta[0] * bunch["xp"]**2))
-            Jy = ((self.ring.optics.local_gamma[1] * bunch["y"]**2) +
-                  (2 * self.ring.optics.local_alpha[1] * bunch["y"] *
-                   bunch["yp"]) +
-                  (self.ring.optics.local_beta[1] * bunch["yp"]**2))
-            phase_advance_x = (
-                2 * np.pi *
-                (self.ring.tune[0] + self.ring.chro[0] * bunch["delta"] +
-                 self.adts_poly[0](Jx) + self.adts_poly[2](Jy)))
-            phase_advance_y = (
-                2 * np.pi *
-                (self.ring.tune[1] + self.ring.chro[1] * bunch["delta"] +
-                 self.adts_poly[1](Jx) + self.adts_poly[3](Jy)))
-
-        # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
-        matrix = np.zeros((6, 6, len(bunch)), dtype=np.float64)
-
-        # Horizontal
-        c_x = np.cos(phase_advance_x)
-        s_x = np.sin(phase_advance_x)
-
-        matrix[0, 0, :] = c_x + self.alpha[0] * s_x
-        matrix[0, 1, :] = self.beta[0] * s_x
-        matrix[0, 2, :] = self.dispersion[0]
-        matrix[1, 0, :] = -1 * self.gamma[0] * s_x
-        matrix[1, 1, :] = c_x - self.alpha[0] * s_x
-        matrix[1, 2, :] = self.dispersion[1]
-        matrix[2, 2, :] = 1
-
-        # Vertical
-        c_y = np.cos(phase_advance_y)
-        s_y = np.sin(phase_advance_y)
-
-        matrix[3, 3, :] = c_y + self.alpha[1] * s_y
-        matrix[3, 4, :] = self.beta[1] * s_y
-        matrix[3, 5, :] = self.dispersion[2]
-        matrix[4, 3, :] = -1 * self.gamma[1] * s_y
-        matrix[4, 4, :] = c_y - self.alpha[1] * s_y
-        matrix[4, 5, :] = self.dispersion[3]
-        matrix[5, 5, :] = 1
-
-        x = (matrix[0, 0] * bunch["x"] + matrix[0, 1] * bunch["xp"] +
-             matrix[0, 2] * bunch["delta"])
-        xp = (matrix[1, 0] * bunch["x"] + matrix[1, 1] * bunch["xp"] +
-              matrix[1, 2] * bunch["delta"])
-        y = (matrix[3, 3] * bunch["y"] + matrix[3, 4] * bunch["yp"] +
-             matrix[3, 5] * bunch["delta"])
-        yp = (matrix[4, 3] * bunch["y"] + matrix[4, 4] * bunch["yp"] +
-              matrix[4, 5] * bunch["delta"])
-
-        bunch["x"] = x
-        bunch["xp"] = xp
-        bunch["y"] = y
-        bunch["yp"] = yp
-
-
 class SkewQuadrupole:
     """
     Thin skew quadrupole element used to introduce betatron coupling (the
@@ -346,6 +245,67 @@ class TransverseMapSector(Element):
         else:
             self.adts_poly = None
 
+    def _compute_chromatic_tune_advances(self, bunch):
+        order = len(self.chro_diff) // 2
+        if order == 1:
+            tune_advance_x = self.chro_diff[0] * bunch["delta"]
+            tune_advance_y = self.chro_diff[1] * bunch["delta"]
+        elif order == 2:
+            tune_advance_x = (self.chro_diff[0] * bunch["delta"] +
+                              self.chro_diff[2] / 2 * bunch["delta"]**2)
+            tune_advance_y = (self.chro_diff[1] * bunch["delta"] +
+                              self.chro_diff[3] / 2 * bunch["delta"]**2)
+        elif order == 3:
+            tune_advance_x = (self.chro_diff[0] * bunch["delta"] +
+                              self.chro_diff[2] / 2 * bunch["delta"]**2 +
+                              self.chro_diff[4] / 6 * bunch["delta"]**3)
+            tune_advance_y = (self.chro_diff[1] * bunch["delta"] +
+                              self.chro_diff[3] / 2 * bunch["delta"]**2 +
+                              self.chro_diff[5] / 6 * bunch["delta"]**3)
+        elif order == 4:
+            tune_advance_x = (self.chro_diff[0] * bunch["delta"] +
+                              self.chro_diff[2] / 2 * bunch["delta"]**2 +
+                              self.chro_diff[4] / 6 * bunch["delta"]**3 +
+                              self.chro_diff[6] / 24 * bunch["delta"]**4)
+            tune_advance_y = (self.chro_diff[1] * bunch["delta"] +
+                              self.chro_diff[3] / 2 * bunch["delta"]**2 +
+                              self.chro_diff[5] / 6 * bunch["delta"]**3 +
+                              self.chro_diff[7] / 24 * bunch["delta"]**4)
+        else:
+            coefs = np.array([1 / factorial(i + 1) for i in range(order + 1)])
+            coefs[0] = 0
+            self.chro_diff = np.concatenate(([0, 0], self.chro_diff))
+            tune_advance_x = np.polynomial.polynomial.Polynomial(
+                self.chro_diff[::2] * coefs)(bunch['delta'])
+            tune_advance_y = np.polynomial.polynomial.Polynomial(
+                self.chro_diff[1::2] * coefs)(bunch['delta'])
+        return tune_advance_x, tune_advance_y
+
+    def _compute_new_coords(self, bunch, tune_advance, plane):
+        if plane == 'x':
+            i, j, coord, mom = 0, 0, 'x', 'xp'
+        elif plane == 'y':
+            i, j, coord, mom = 1, 2, 'y', 'yp'
+        else:
+            raise ValueError('plane should be either x or y')
+        c_u = np.cos(2 * np.pi * tune_advance)
+        s_u = np.sin(2 * np.pi * tune_advance)
+        M00 = np.sqrt(
+            self.beta1[i] / self.beta0[i]) * (c_u + self.alpha0[i] * s_u)
+        M01 = np.sqrt(self.beta0[i] * self.beta1[i]) * s_u
+        M02 = (self.dispersion1[j] - M00 * self.dispersion0[j] -
+               M01 * self.dispersion0[j + 1])
+        M10 = ((self.alpha0[i] - self.alpha1[i]) * c_u -
+               (1 + self.alpha0[i] * self.alpha1[i]) * s_u) / np.sqrt(
+                   self.beta0[i] * self.beta1[i])
+        M11 = np.sqrt(
+            self.beta0[i] / self.beta1[i]) * (c_u - self.alpha1[i] * s_u)
+        M12 = (self.dispersion1[j + 1] - M10 * self.dispersion0[j] -
+               M11 * self.dispersion0[j + 1])
+        u = (M00 * bunch[coord] + M01 * bunch[mom] + M02 * bunch["delta"])
+        up = (M10 * bunch[coord] + M11 * bunch[mom] + M12 * bunch["delta"])
+        return u, up
+
     @Element.parallel
     def track(self, bunch):
         """
@@ -357,85 +317,45 @@ class TransverseMapSector(Element):
         ----------
         bunch : Bunch or Beam object
         """
-
-        # Compute phase advance which depends on energy via chromaticity and ADTS
-        if self.adts_poly is None:
-            phase_advance_x = (
-                2 * np.pi *
-                (self.tune_diff[0] + self.chro_diff[0] * bunch["delta"]))
-            phase_advance_y = (
-                2 * np.pi *
-                (self.tune_diff[1] + self.chro_diff[1] * bunch["delta"]))
-        else:
+        tune_advance_x = self.tune_diff[0]
+        tune_advance_y = self.tune_diff[1]
+        # Compute tune advance which depends on energy via chromaticity and ADTS
+        if (np.array(self.chro_diff) != 0).any():
+            tune_advance_x_chro, tune_advance_y_chro = self._compute_chromatic_tune_advances(
+                bunch)
+            tune_advance_x += tune_advance_x_chro
+            tune_advance_y += tune_advance_y_chro
+
+        if self.adts_poly is not None:
             Jx = ((self.gamma0[0] * bunch["x"]**2) +
                   (2 * self.alpha0[0] * bunch["x"] * self["xp"]) +
                   (self.beta0[0] * bunch["xp"]**2))
             Jy = ((self.gamma0[1] * bunch["y"]**2) +
                   (2 * self.alpha0[1] * bunch["y"] * bunch["yp"]) +
                   (self.beta0[1] * bunch["yp"]**2))
-            phase_advance_x = (
-                2 * np.pi *
-                (self.tune_diff[0] + self.chro_diff[0] * bunch["delta"] +
-                 self.adts_poly[0](Jx) + self.adts_poly[2](Jy)))
-            phase_advance_y = (
-                2 * np.pi *
-                (self.tune_diff[1] + self.chro_diff[1] * bunch["delta"] +
-                 self.adts_poly[1](Jx) + self.adts_poly[3](Jy)))
-
-        # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
-        matrix = np.zeros((6, 6, len(bunch)))
-
-        # Horizontal
-        matrix[0, 0, :] = np.sqrt(self.beta1[0] / self.beta0[0]) * (
-            np.cos(phase_advance_x) + self.alpha0[0] * np.sin(phase_advance_x))
-        matrix[0, 1, :] = np.sqrt(
-            self.beta0[0] * self.beta1[0]) * np.sin(phase_advance_x)
-        matrix[0, 2, :] = (self.dispersion1[0] -
-                           matrix[0, 0, :] * self.dispersion0[0] -
-                           matrix[0, 1, :] * self.dispersion0[1])
-        matrix[1, 0, :] = (
-            (self.alpha0[0] - self.alpha1[0]) * np.cos(phase_advance_x) -
-            (1 + self.alpha0[0] * self.alpha1[0]) *
-            np.sin(phase_advance_x)) / np.sqrt(self.beta0[0] * self.beta1[0])
-        matrix[1, 1, :] = np.sqrt(self.beta0[0] / self.beta1[0]) * (
-            np.cos(phase_advance_x) - self.alpha1[0] * np.sin(phase_advance_x))
-        matrix[1, 2, :] = (self.dispersion1[1] -
-                           matrix[1, 0, :] * self.dispersion0[0] -
-                           matrix[1, 1, :] * self.dispersion0[1])
-        matrix[2, 2, :] = 1
-
-        # Vertical
-        matrix[3, 3, :] = np.sqrt(self.beta1[1] / self.beta0[1]) * (
-            np.cos(phase_advance_y) + self.alpha0[1] * np.sin(phase_advance_y))
-        matrix[3, 4, :] = np.sqrt(
-            self.beta0[1] * self.beta1[1]) * np.sin(phase_advance_y)
-        matrix[3, 5, :] = (self.dispersion1[2] -
-                           matrix[3, 3, :] * self.dispersion0[2] -
-                           matrix[3, 4, :] * self.dispersion0[3])
-        matrix[4, 3, :] = (
-            (self.alpha0[1] - self.alpha1[1]) * np.cos(phase_advance_y) -
-            (1 + self.alpha0[1] * self.alpha1[1]) *
-            np.sin(phase_advance_y)) / np.sqrt(self.beta0[1] * self.beta1[1])
-        matrix[4, 4, :] = np.sqrt(self.beta0[1] / self.beta1[1]) * (
-            np.cos(phase_advance_y) - self.alpha1[1] * np.sin(phase_advance_y))
-        matrix[4, 5, :] = (self.dispersion1[3] -
-                           matrix[4, 3, :] * self.dispersion0[2] -
-                           matrix[4, 4, :] * self.dispersion0[3])
-        matrix[5, 5, :] = 1
-
-        x = (matrix[0, 0, :] * bunch["x"] + matrix[0, 1, :] * bunch["xp"] +
-             matrix[0, 2, :] * bunch["delta"])
-        xp = (matrix[1, 0, :] * bunch["x"] + matrix[1, 1, :] * bunch["xp"] +
-              matrix[1, 2, :] * bunch["delta"])
-        y = (matrix[3, 3, :] * bunch["y"] + matrix[3, 4, :] * bunch["yp"] +
-             matrix[3, 5, :] * bunch["delta"])
-        yp = (matrix[4, 3, :] * bunch["y"] + matrix[4, 4, :] * bunch["yp"] +
-              matrix[4, 5, :] * bunch["delta"])
-
-        bunch["x"] = x
-        bunch["xp"] = xp
-        bunch["y"] = y
-        bunch["yp"] = yp
+            tune_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
+            tune_advance_x += (self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
+
+        bunch['x'], bunch['xp'] = self._compute_new_coords(
+            bunch, tune_advance_x, 'x')
+        bunch['y'], bunch['yp'] = self._compute_new_coords(
+            bunch, tune_advance_y, 'y')
+
+
+class TransverseMap(TransverseMapSector):
+    """
+    Transverse map for a single turn in the synchrotron.
+
+    Parameters
+    ----------
+    ring : Synchrotron object
+    """
+
+    def __init__(self, ring):
+        super().__init__(ring, ring.optics.local_alpha, ring.optics.local_beta,
+                         ring.optics.local_dispersion, ring.optics.local_alpha,
+                         ring.optics.local_beta, ring.optics.local_dispersion,
+                         2 * np.pi * ring.tune, ring.chro, ring.adts)
 
 
 def transverse_map_sector_generator(ring, positions):
@@ -470,36 +390,29 @@ def transverse_map_sector_generator(ring, positions):
     if ring.optics.use_local_values:
         for i in range(N_sec):
             sectors.append(
-                TransverseMapSector(
-                    ring,
-                    ring.optics.local_alpha,
-                    ring.optics.local_beta,
-                    ring.optics.local_dispersion,
-                    ring.optics.local_alpha,
-                    ring.optics.local_beta,
-                    ring.optics.local_dispersion,
-                    ring.tune / N_sec,
-                    ring.chro / N_sec,
-                ))
+                TransverseMapSector(ring,
+                                    ring.optics.local_alpha,
+                                    ring.optics.local_beta,
+                                    ring.optics.local_dispersion,
+                                    ring.optics.local_alpha,
+                                    ring.optics.local_beta,
+                                    ring.optics.local_dispersion,
+                                    2 * np.pi * ring.tune / N_sec,
+                                    ring.chro / N_sec,
+                                    adts=ring.adts /
+                                    N_sec if ring.adts else None))
     else:
         import at
 
-        def _compute_chro(ring, pos, dp=1e-4):
+        def _compute_chro(ring, pos, dp=1e-2, 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=dp,
+                                                               n_points=100,
+                                                               order=order)
+            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 15235ecd1b60f38a4984074ab71c87c1fba54a85..41458eff38988e909c1f3c90b2d4863c0a3e12ae 100644
--- a/mbtrack2/tracking/synchrotron.py
+++ b/mbtrack2/tracking/synchrotron.py
@@ -109,6 +109,9 @@ class Synchrotron:
     get_adts()
         Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar
         componenet from AT lattice.
+    get_chroma()
+        Compute chromaticity (linear and nonlinear) from AT lattice and 
+        update the property.
     get_mcf_order()
         Compute momentum compaction factor up to 3rd order from AT lattice.
     get_longitudinal_twiss(V)
@@ -387,6 +390,25 @@ class Synchrotron:
         coef_yy = np.array([det.A3 / 2, 0])
         self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]
 
+    def get_chroma(self, order=4, dpm=0.02, n_points=100):
+        """
+        Compute chromaticity (linear and nonlinear) from AT lattice and update the property.
+        """
+        import at
+        if self.optics.use_local_values:
+            raise ValueError("Value needs to be provided manualy as no AT" +
+                             " lattice file is loaded.")
+        fit, dpa, tune = at.physics.nonlinear.chromaticity(self.optics.lattice,
+                                                           method='linopt',
+                                                           dpm=dpm,
+                                                           n_points=n_points,
+                                                           order=order)
+        chrox, chroy = fit
+        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):
         """
         Compute momentum compaction factor up to 3rd order from AT lattice.
@@ -408,7 +430,7 @@ class Synchrotron:
         """
         import at
         if self.optics.use_local_values:
-            raise ValueError("ADTS needs to be provided manualy as no AT" +
+            raise ValueError("Value needs to be provided manualy as no AT" +
                              " lattice file is loaded.")
         deltamin = -1e-4
         deltamax = 1e-4