Nonlinear chromaticity

Hello,

Nonlinear chromaticity implementation is now ready for a merge request. The list of changes is below. To trigger nonlinear chromaticity, you simply need to define the ring.chro = [Qx', Qy', Qx'', Qy'',..]. All orders are implemented. Orders < 4 are implemented separately to gain some performance over the numpy.polynomial.

Jupyter notebook to test the implementation: Nonlinear_chromaticity.ipynb

• added nonlinear chromaticity calculation to TransverseMap and TransverseSectorMap
• TransverseMap now inherits from TransverseSectorMap
• chromaticity phase shift is computed in _compute_chromatic_phase_advances()
• some optimisation of the code, (250 ms -> <50 ms) speed up for typical simulation Slight improvements for SynchrotronRadiation
• np.random.normal -> np.random.standard_normal, 10% faster
• N = len(bunch) to compute the number of particles only once

mbtrack2/tracking/element.py

@@ -404,22 +404,15 @@ def transverse_map_sector_generator(ring, positions):
     else:
         import at
 
-        def _compute_chro(ring, pos, dp=1e-4):
+        def _compute_chro(ring, pos, dp=1e-4, 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=0.02,
+                                                               n_points=100,
+                                                               order=order)
+            Chrox, Chroy = fit
             chrox = np.interp(pos, s, Chrox)
             chroy = np.interp(pos, s, Chroy)