diff --git a/.gitignore b/.gitignore index 292503b26..4495f13d0 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,5 @@ xtrack/prebuilt_kernels !xtrack/prebuilt_kernels/*.py checkpoint_restart.dat +.__afs* +*.c diff --git a/tests/test_coupled_line_segment_map.py b/tests/test_coupled_line_segment_map.py new file mode 100644 index 000000000..e00d2137e --- /dev/null +++ b/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)) + + \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 97158c560..cc34f3a93 100644 --- a/xtrack/beam_elements/elements.py +++ b/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 diff --git a/xtrack/beam_elements/elements_src/linesegmentmap.h b/xtrack/beam_elements/elements_src/linesegmentmap.h index db6d3b5f3..8dc2d8e9c 100644 --- a/xtrack/beam_elements/elements_src/linesegmentmap.h +++ b/xtrack/beam_elements/elements_src/linesegmentmap.h @@ -101,6 +101,19 @@ void transverse_motion(LocalParticle *part0, LineSegmentMapData el){ int64_t const ndqx = LineSegmentMapData_len_coeffs_dqx(el); int64_t const ndqy = LineSegmentMapData_len_coeffs_dqy(el); + double const c11_0= LineSegmentMapData_get_c11(el,0); + double const c12_0= LineSegmentMapData_get_c12(el,0); + double const c21_0= LineSegmentMapData_get_c21(el,0); + double const c22_0= LineSegmentMapData_get_c22(el,0); + + double const c11_1= LineSegmentMapData_get_c11(el,1);; + double const c12_1= LineSegmentMapData_get_c12(el,1);; + double const c21_1= LineSegmentMapData_get_c21(el,1);; + double const c22_1= LineSegmentMapData_get_c22(el,1); + + double const coupling_0 = c11_0 + c22_0 + c12_0 + c21_0; + double const coupling_1 = c11_1 + c22_1 + c12_1 + c21_1; + int64_t detuning; double sin_x = 0; double cos_x = 0; @@ -180,18 +193,136 @@ void transverse_motion(LocalParticle *part0, LineSegmentMapData el){ )/sqrt_betprod_y; double const M11_y = (cos_y-alfy_1*sin_y)/sqrt_betratio_y; - double const x_out = M00_x*LocalParticle_get_x(part) + M01_x * LocalParticle_get_px(part); - double const px_out = M10_x*LocalParticle_get_x(part) + M11_x * LocalParticle_get_px(part); - double const y_out = M00_y*LocalParticle_get_y(part) + M01_y * LocalParticle_get_py(part); - double const py_out = M10_y*LocalParticle_get_y(part) + M11_y * LocalParticle_get_py(part); + if (coupling_0 == 0.0 && coupling_1 == 0.0) { + // No coupling matrices, use simple transfer matrix + double const x_out = M00_x*LocalParticle_get_x(part) + M01_x * LocalParticle_get_px(part); + double const px_out = M10_x*LocalParticle_get_x(part) + M11_x * LocalParticle_get_px(part); + double const y_out = M00_y*LocalParticle_get_y(part) + M01_y * LocalParticle_get_py(part); + double const py_out = M10_y*LocalParticle_get_y(part) + M11_y * LocalParticle_get_py(part); + + LocalParticle_set_x(part, x_out); + LocalParticle_set_px(part, px_out); + LocalParticle_set_y(part, y_out); + LocalParticle_set_py(part, py_out); + } - LocalParticle_set_x(part, x_out); - LocalParticle_set_px(part, px_out); - LocalParticle_set_y(part, y_out); - LocalParticle_set_py(part, py_out); - END_PER_PARTICLE_BLOCK; + /* Construct 4x4 symplectic matrix R and its inverse from a 2x2 matrix C */ + void compute_matrices_V(double const c11, double const c12, + double const c21, double const c22, + double const gamma, + double V[16]) + { + /* Initialize R and Rinv to zero */ + for(int i=0;i<16;i++) { + V[i] = 0.0; + } + /* Identity 2x2 blocks (top-left and bottom-right) */ + V[0] = gamma; V[5] = gamma; + V[10] = gamma; V[15] = gamma; + /* C block (top-right) */ + V[2] = c11; V[3] = c12; + V[6] = c21; V[7] = c22; + /* -C^inv block (bottom-left) */ + V[8] = -c22; V[9] = c12; + V[12] = c21; V[13] = -c11; + } + + /* Construct 4x4 symplectic matrix R and its inverse from a 2x2 matrix C */ + void compute_matrices_Vinv(double const c11, double const c12, + double const c21, double const c22, + double const gamma, + double V_inv[16]) + { + /* Initialize R and Rinv to zero */ + for(int i=0;i<16;i++) { + V_inv[i] = 0.0; + } + /* Identity 2x2 blocks (top-left and bottom-right) */ + V_inv[0] = gamma; V_inv[5] = gamma; + V_inv[10] = gamma; V_inv[15] = gamma; + /* -C block (top-right) */ + V_inv[2] = -c11; V_inv[3] = -c12; + V_inv[6] = -c21; V_inv[7] = -c22; + /*C^inv block (bottom-left) */ + V_inv[8] = c22; V_inv[9] = -c12; + V_inv[12] = -c21; V_inv[13] = c11; + } + + + if (coupling_0 != 0.0 || coupling_1 != 0.0) { + + double const det_0 = c11_0 * c22_0 - c12_0 * c21_0; + double const det_1 = c11_1 * c22_1 - c12_1 * c21_1; + + double const gamma_0 = sqrt(1 - det_0); + double const gamma_1 = sqrt(1 - det_1); + + double V_1[16], V_0inv[16]; + double M[16]; + double T[16]; + double temp[16]; + + compute_matrices_V(c11_1, c12_1, c21_1, c22_1, gamma_1, V_1); + compute_matrices_Vinv(c11_0, c12_0, c21_0, c22_0, gamma_0, V_0inv); + + //the final transformation is V_1 * M * V_0inv + + for(int i=0;i<16;i++) { + M[i] = 0.0; + } + M[0] = M00_x; M[1] = M01_x; M[2]=0.0; M[3]=0.0; + M[4] = M10_x; M[5] = M11_x; M[6]=0.0; M[7]=0.0; + M[8] = 0.0; M[9] = 0.0; M[10]=M00_y; M[11]=M01_y; + M[12]= 0.0; M[13]= 0.0; M[14]=M10_y; M[15]=M11_y; + + for(int i=0;i<16;i++) { + T[i] = 0.0; + } + + for(int i=0;i<16;i++) { + temp[i] = 0.0; + } + + // temp = M * V_0inv + for (int i=0;i<4;i++) { + for (int j=0;j<4;j++) { + temp[i*4+j] = 0.0; + for (int k=0;k<4;k++) + temp[i*4+j] += M[i*4+k]*V_0inv[k*4+j]; + } + } + // T = V_1 * temp + for (int i=0;i<4;i++) { + for (int j=0;j<4;j++) { + T[i*4+j] = 0.0; + for (int k=0;k<4;k++) + T[i*4+j] += V_1[i*4+k]*temp[k*4+j]; + } + } + + double const old_x = LocalParticle_get_x(part); + double const old_px = LocalParticle_get_px(part); + double const old_y = LocalParticle_get_y(part); + double const old_py = LocalParticle_get_py(part); + + double const x_out = ( + T[0]*old_x + T[1]*old_px + T[2]*old_y + T[3]*old_py); + double const px_out = ( + T[4]*old_x + T[5]*old_px + T[6]*old_y + T[7]*old_py); + double const y_out = ( + T[8]*old_x + T[9]*old_px + T[10]*old_y + T[11]*old_py); + double const py_out = ( + T[12]*old_x + T[13]*old_px + T[14]*old_y + T[15]*old_py); + + LocalParticle_set_x(part, x_out); + LocalParticle_set_px(part, px_out); + LocalParticle_set_y(part, y_out); + LocalParticle_set_py(part, py_out); + END_PER_PARTICLE_BLOCK; + } } + GPUFUN void longitudinal_motion(LocalParticle *part0, LineSegmentMapData el){