Merge remote-tracking branch 'ColasDroin/feature/update_dispersion' into release/v0.52.0

Author: szymonlopaciuk

tests/test_twiss.py

@@ -1177,6 +1177,30 @@ def test_crab_dispersion(test_context):
     assert np.allclose(tw6d_rf_on['dx_zeta'], tw4d_rf_off['dx_zeta'], rtol=0, atol=1e-7)
     assert np.allclose(tw6d_rf_on['dy_zeta'], tw4d_rf_off['dy_zeta'], rtol=0, atol=1e-7)
 
+@for_all_test_contexts
+def test_higher_order_crab_dispersion(test_context):
+
+    collider = xt.Multiline.from_json(test_data_folder /
+                        'hllhc15_collider/collider_00_from_mad.json')
+    collider.build_trackers(_context=test_context)
+
+    collider.vars['vrf400'] = 16
+    collider.vars['on_crab1'] = -190
+    collider.vars['on_crab5'] = -190
+
+    line = collider.lhcb1
+
+    tw6d = line.twiss(method='6d')
+    dz = 1e-3
+    tw4d_plus = line.twiss(method='4d', zeta0=dz, freeze_longitudinal=True)
+    tw4d_minus = line.twiss(method='4d', zeta0=-dz, freeze_longitudinal=True)
+    dpx_dzeta_expected = (tw4d_plus.px -  tw4d_minus.px) / (2*dz)
+    dpy_dzeta_expected = (tw4d_plus.py -  tw4d_minus.py) / (2*dz)
+
+    assert np.allclose(tw6d['dpx_zeta'], dpx_dzeta_expected, rtol=0, atol=1e-7)
+    assert np.allclose(tw6d['dpy_zeta'], dpy_dzeta_expected, rtol=0, atol=1e-7)
+
+
 @for_all_test_contexts
 def test_twiss_group_compounds(test_context):
 

xtrack/twiss.py

@@ -181,16 +181,16 @@ def twiss_line(line, particle_ref=None, method=None,
             - dx: horizontal dispersion (d x / d delta) in meters
             - dy: vertical dispersion (d y / d delta) in meters
             - dzeta: longitudinal dispersion (d zeta / d delta) in meters
-            - dpx: horizontal dispersion (d px / d delta)
-            - dpy: vertical dispersion (d y / d delta)
+            - dpx: higher order horizontal dispersion (d px / d delta)
+            - dpy: higher order vertical dispersion (d py / d delta)
             - ddx: horizontal second order dispersion (d^2 x / d delta^2) in meters
             - ddy: vertical second order dispersion (d^2 y / d delta^2) in meters
             - ddpx: horizontal second order dispersion (d^2 px / d delta^2)
             - ddpy: vertical second order dispersion (d^2 py / d delta^2)
             - dx_zeta: horizontal crab dispersion (d x / d zeta)
             - dy_zeta: vertical crab dispersion (d y / d zeta)
-            - dpx_zeta: horizontal crab dispersion (d px / d zeta)
-            - dpy_zeta: vertical crab dispersion (d py / d zeta)
+            - dpx_zeta: higher order horizontal crab dispersion (d px / d zeta)
+            - dpy_zeta: higher order vertical crab dispersion (d py / d zeta)
             - ax_chrom: chromatic function (d alfx / d delta - alfx / betx d betx / d delta)
             - ay_chrom: chromatic function (d alfy / d delta - alfy / bety d bety / d delta)
             - bx_chrom: chromatic function (d betx / d delta)
@@ -1051,11 +1051,17 @@ def _compute_lattice_functions(Ws, use_full_inverse, s_co):
     muy = np.unwrap(temp_phiy) / 2  /np.pi
     muzeta = np.unwrap(phizeta) / 2 / np.pi
 
+    # Crab dispersion
     dx_zeta = (Ws[:, 0, 4] - Ws[:, 0, 5] * Ws[:, 5, 4] / Ws[:, 5, 5]) / (
                Ws[:, 4, 4] - Ws[:, 4, 5] * Ws[:, 5, 4] / Ws[:, 5, 5])
+    dpx_zeta = (Ws[:, 1, 4] - Ws[:, 1, 5] * Ws[:, 5, 4] / Ws[:, 5, 5]) / (
+                Ws[:, 4, 4] - Ws[:, 4, 5] * Ws[:, 5, 4] / Ws[:, 5, 5])
     dy_zeta = (Ws[:, 2, 4] - Ws[:, 2, 5] * Ws[:, 5, 4] / Ws[:, 5, 5]) / (
                 Ws[:, 4, 4] - Ws[:, 4, 5] * Ws[:, 5, 4] / Ws[:, 5, 5])
+    dpy_zeta = (Ws[:, 3, 4] - Ws[:, 3, 5] * Ws[:, 5, 4] / Ws[:, 5, 5]) / (
+                Ws[:, 4, 4] - Ws[:, 4, 5] * Ws[:, 5, 4] / Ws[:, 5, 5])
 
+    # Dispersion
     dx_pzeta = (Ws[:, 0, 5] - Ws[:, 0, 4] * Ws[:, 4, 5] / Ws[:, 4, 4]) / (
                 Ws[:, 5, 5] - Ws[:, 5, 4] * Ws[:, 4, 5] / Ws[:, 4, 4])
     dpx_pzeta = (Ws[:, 1, 5] - Ws[:, 1, 4] * Ws[:, 4, 5] / Ws[:, 4, 4]) / (
@@ -1081,7 +1087,9 @@ def _compute_lattice_functions(Ws, use_full_inverse, s_co):
         'dy': dy_pzeta,
         'dpy': dpy_pzeta,
         'dx_zeta': dx_zeta,
+        'dpx_zeta': dpx_zeta,
         'dy_zeta': dy_zeta,
+        'dpy_zeta': dpy_zeta,
         'betx1': betx1,
         'bety1': bety1,
         'betx2': betx2,