Add coupled transverse motion to LineSegmentMap with Cij matrix

.gitignore

@@ -16,3 +16,5 @@ xtrack/prebuilt_kernels
 !xtrack/prebuilt_kernels/*.py
 checkpoint_restart.dat
 
+.__afs*
+*.c

tests/test_coupled_line_segment_map.py

@@ -0,0 +1,203 @@
+import numpy as np
+import xobjects as xo
+import xtrack as xt
+import xpart as xp
+
+
+# ============================================================
+# RANDOM TEST SETTINGS
+# ============================================================
+
+N_TESTS = 5
+RDT_MAX = 1e-2
+
+
+# ============================================================
+# FCC PARAMETERS (FROM YOU)
+# ============================================================
+
+energy = 120
+p0c = 120e9
+n_ips = 4
+
+beta_x = 240e-3
+beta_y = 1.0e-3
+
+alpha_x = 0.0
+alpha_y = 0.0
+
+Qx = 398.150 / n_ips
+Qy = 398.220 / n_ips
+Qs = 0.0334 / n_ips
+
+phi = 15e-3
+k2_factor = 0.50
+
+sigma_delta_bs = 0.00176
+sigma_z_bs = 5.59e-3
+beta_s_bs = sigma_z_bs / sigma_delta_bs
+
+
+# sextupole optics (can tune later)
+beta_x_sext = 3.0
+beta_y_sext = 500.0
+
+alpha_x_sext = 0.0
+alpha_y_sext = 0.0
+
+
+context = xo.ContextCpu()
+
+
+# ============================================================
+# RDT → Cbar
+# ============================================================
+
+#D. Sagan and D. Rubin. Linear analysis of coupled lattices.
+#$Phys. Rev. ST Accel. Beams, 2:074001, Jul 1999.
+
+#R. Calaga, R. Tomás, and A. Franchi. Betatron coupling:
+#Merging Hamiltonian and matrix approaches. Phys. Rev. ST
+#Accel. Beams, 8:034001, Mar 2005
+
+#Vaibhavi Gawas, Frank Zimmermann, Rogelio Tomas Garcia,
+#and Xavier Buffat. MODELLING LINEAR COUPLING
+#FOR FCC-ee IN XSUITE . In Proc. IPAC’26.
+
+
+def rdts_to_Cbar(f1001, f1010):
+
+    re1010, im1010 = np.real(f1010), np.imag(f1010)
+    re1001, im1001 = np.real(f1001), np.imag(f1001)
+
+    S11 = (im1010 + im1001)
+    S22 = (im1001 - im1010)
+    S12 = (-re1010 + re1001)
+    S21 = (-re1010 - re1001)
+
+    S = np.array([[S11, S12], [S21, S22]])
+
+    detS = np.linalg.det(S)
+    gamma = 1.0 / np.sqrt(1.0 + 4.0 * detS)
+
+    return 2 * gamma * S
+
+
+def Cbar_to_Cphys(Cbar):
+
+    C11 = Cbar[0, 0] * np.sqrt(beta_x / beta_y)
+    C12 = Cbar[0, 1] * np.sqrt(beta_x * beta_y)
+    C21 = Cbar[1, 0] / np.sqrt(beta_x * beta_y)
+    C22 = Cbar[1, 1] * np.sqrt(beta_y / beta_x)
+
+    return C11, C12, C21, C22
+
+
+# ============================================================
+# TEST LOOP
+# ============================================================
+
+for i in range(N_TESTS):
+
+    # -------- random RDTs ----------
+    r1001 = np.random.uniform(0, RDT_MAX)
+    i1001 = np.random.uniform(0, RDT_MAX)
+    r1010 = np.random.uniform(0, RDT_MAX)
+    i1010 = np.random.uniform(0, RDT_MAX)
+
+    f1001 = r1001 + 1j * i1001
+    f1010 = r1010 + 1j * i1010
+
+    Cbar = rdts_to_Cbar(f1001, f1010)
+    c11, c12, c21, c22 = Cbar_to_Cphys(Cbar)
+
+    # -------- sext strengths ----------
+    k2_left = k2_factor / (2 * phi * beta_y * beta_y_sext) * np.sqrt(beta_x / beta_x_sext)
+    k2_right = k2_left
+
+    # -------- arcs ----------
+    arc_left = xt.LineSegmentMap(
+        _context=context,
+        qx=0.0,
+        qy=0.25,
+        qs=0.0,
+        betx=[beta_x, beta_x_sext],
+        bety=[beta_y, beta_y_sext],
+        alfx=[alpha_x, alpha_x_sext],
+        alfy=[alpha_y, alpha_y_sext],
+        c11=[c11, 0.0],
+        c12=[c12, 0.0],
+        c21=[c21, 0.0],
+        c22=[c22, 0.0],
+        bets=beta_s_bs,
+    )
+
+    arc_mid = xt.LineSegmentMap(
+        _context=context,
+        qx=Qx,
+        qy=Qy - 0.5,
+        qs=Qs,
+        betx=[beta_x_sext, beta_x_sext],
+        bety=[beta_y_sext, beta_y_sext],
+        alfx=[alpha_x_sext, alpha_x_sext],
+        alfy=[alpha_y_sext, alpha_y_sext],
+        bets=beta_s_bs,
+    )
+
+    arc_right = xt.LineSegmentMap(
+        _context=context,
+        qx=0.0,
+        qy=0.25,
+        qs=0.0,
+        betx=[beta_x_sext, beta_x],
+        bety=[beta_y_sext, beta_y],
+        alfx=[alpha_x_sext, alpha_x],
+        alfy=[alpha_y_sext, alpha_y],
+        c11=[0.0, c11],
+        c12=[0.0, c12],
+        c21=[0.0, c21],
+        c22=[0.0, c22],
+        bets=beta_s_bs,
+    )
+
+    sext_left = xt.Multipole(order=2, knl=[0, 0, k2_left], _context=context)
+    sext_right = xt.Multipole(order=2, knl=[0, 0, -k2_right], _context=context)
+
+    line = xt.Line(elements=[
+        arc_mid,
+        sext_right,
+        arc_right,
+        arc_left,
+        sext_left
+    ])
+
+    line.particle_ref = xp.Particles(
+        _context=context,
+        p0c=p0c,
+        mass0=xp.ELECTRON_MASS_EV
+    )
+
+    line.build_tracker(_context = context, use_prebuilt_kernels = False)
+
+    tw = line.twiss(method="4d", coupling_edw_teng=True)
+
+    f1001_tw = tw.f1001[3]
+    f1010_tw = tw.f1010[3]
+
+    print("\n==============================")
+    print("TEST", i + 1)
+
+    print("INPUT")
+    print("f1001 =", f1001)
+    print("f1010 =", f1010)
+
+    print("\nTWISS")
+    print("f1001 =", f1001_tw)
+    print("f1010 =", f1010_tw)
+
+    ## Xsuite twiss.py outputs -np.conj(rdt) of the rdt which should match the input
+    print("\nnegative of Conjugate")
+    print("f1001 =",-np.conj(f1001_tw))
+    print("f1010 =",-np.conj(f1010_tw))
+
+

xtrack/beam_elements/elements.py

@@ -3619,6 +3619,11 @@ class LineSegmentMap(BeamElement):
         'alfx': xo.Float64[2],
         'alfy': xo.Float64[2],
 
+        'c11': xo.Float64[2],
+        'c12': xo.Float64[2],
+        'c21': xo.Float64[2],
+        'c22': xo.Float64[2],
+
         'dx': xo.Float64[2],
         'dpx': xo.Float64[2],
         'dy': xo.Float64[2],
@@ -3665,6 +3670,7 @@ class LineSegmentMap(BeamElement):
 
     def __init__(self, length=0., qx=0, qy=0,
             betx=1., bety=1., alfx=0., alfy=0.,
+            c11=0., c12=0., c21=0., c22=0.,
             dx=0., dpx=0., dy=0., dpy=0.,
             x_ref=0.0, px_ref=0.0, y_ref=0.0, py_ref=0.0,
             longitudinal_mode=None,
@@ -3707,6 +3713,18 @@ def __init__(self, length=0., qx=0, qy=0,
         alfy : tuple of length 2 or float
             Vertical alpha function at the entrance and exit of the segment.
             If a float is given, the same value is used for both entrance and exit.
+        c11 : tuple of length 2 or float
+            Coupling matrix element
+            If a float is given, the same value is used for both entrance and exit
+        c12 : tuple of length 2 or float
+            Coupling matrix element
+            If a float is given, the same value is used for both entrance and exit
+        c21 : tuple of length 2 or float
+            Coupling matrix element
+            If a float is given, the same value is used for both entrance and exit
+        c22 : tuple of length 2 or float
+            Coupling matrix element
+            If a float is given, the same value is used for both entrance and exit
         dx : tuple of length 2 or float
             Horizontal dispersion at the entrance and exit of the segment.
             If a float is given, the same value is used for both entrance and exit.
@@ -3963,6 +3981,18 @@ def __init__(self, length=0., qx=0, qy=0,
         if np.isscalar(alfy): alfy = [alfy, alfy]
         else: assert len(alfy) == 2
 
+        if np.isscalar(c11): c11 = [c11, c11]
+        else: assert len(c11) == 2    
+
+        if np.isscalar(c12): c12 = [c12, c12]
+        else: assert len(c12) == 2
+
+        if np.isscalar(c21): c21 = [c21, c21]
+        else: assert len(c21) == 2
+
+        if np.isscalar(c22): c22 = [c22, c22]
+        else: assert len(c22) == 2
+
         if np.isscalar(dx): dx = [dx, dx]
         else: assert len(dx) == 2
 
@@ -3991,6 +4021,10 @@ def __init__(self, length=0., qx=0, qy=0,
         nargs['bety'] = bety
         nargs['alfx'] = alfx
         nargs['alfy'] = alfy
+        nargs['c11'] = c11
+        nargs['c12'] = c12
+        nargs['c21'] = c21
+        nargs['c22'] = c22
         nargs['dx'] = dx
         nargs['dpx'] = dpx
         nargs['dy'] = dy