xtrack 0.98.12__tar.gz → 0.98.14__tar.gz

{xtrack-0.98.12/xtrack.egg-info → xtrack-0.98.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xtrack
-Version: 0.98.12
+Version: 0.98.14
 Summary: Tracking library for particle accelerators
 Author: G. Iadarola et al.
 License-Expression: Apache-2.0

{xtrack-0.98.12 → xtrack-0.98.14}/tests/test_coupling_edwards_teng.py

@@ -69,24 +69,23 @@ def test_coupling_edwards_teng():
                             betx_mad_at_s, bety_mad_at_s, alfx_mad_at_s, alfy_mad_at_s)
 
     xo.assert_allclose(tw.r11_edw_teng, r11_mad_at_s,
-                    rtol=1e-5, atol=1e-5*np.max(np.abs(r11_mad_at_s)))
+                    rtol=1e-5, atol=5e-4*np.max(np.abs(r11_mad_at_s)))
     xo.assert_allclose(tw.r12_edw_teng, r12_mad_at_s,
-                    rtol=1e-5, atol=1e-5*np.max(np.abs(r12_mad_at_s)))
+                    rtol=1e-5, atol=5e-4*np.max(np.abs(r12_mad_at_s)))
     xo.assert_allclose(tw.r21_edw_teng, r21_mad_at_s,
-                    rtol=1e-5, atol=1e-5*np.max(np.abs(r21_mad_at_s)))
+                    rtol=1e-5, atol=6e-4*np.max(np.abs(r21_mad_at_s)))
     xo.assert_allclose(tw.r22_edw_teng, r22_mad_at_s,
-                    rtol=1e-5, atol=1e-5*np.max(np.abs(r22_mad_at_s)))
-    xo.assert_allclose(tw.betx_edw_teng, betx_mad_at_s, atol=0, rtol=1e-5)
-
-    xo.assert_allclose(tw.betx_edw_teng, betx_mad_at_s, atol=0, rtol=5e-8)
-    xo.assert_allclose(tw.alfx_edw_teng, alfx_mad_at_s, atol=1e-4, rtol=1e-8)
-    xo.assert_allclose(tw.bety_edw_teng, bety_mad_at_s, atol=0, rtol=5e-8)
-    xo.assert_allclose(tw.alfy_edw_teng, alfy_mad_at_s, atol=1e-4, rtol=1e-8)
-
-    xo.assert_allclose(tw.f1001, rdt_mad_at_s['f1001'],
-                    atol=1e-5 * np.max(np.abs(rdt_mad_at_s['f1001'])))
-    xo.assert_allclose(tw.f1010, rdt_mad_at_s['f1010'],
-                    atol=1e-5 * np.max(np.abs(rdt_mad_at_s['f1010'])))
+                    rtol=1e-5, atol=5e-4*np.max(np.abs(r22_mad_at_s)))
+    xo.assert_allclose(tw.betx_edw_teng, betx_mad_at_s, atol=0, rtol=1e-6)
+    xo.assert_allclose(tw.alfx_edw_teng, alfx_mad_at_s, atol=1e-4, rtol=1e-6)
+    xo.assert_allclose(tw.bety_edw_teng, bety_mad_at_s, atol=0, rtol=1e-6)
+    xo.assert_allclose(tw.alfy_edw_teng, alfy_mad_at_s, atol=1e-4, rtol=1e-6)
+
+    xo.assert_allclose(tw.f1001, rdt_mad_at_s['f1001'], rtol=1e-3,
+                    atol=5e-4 * np.max(np.abs(rdt_mad_at_s['f1001'])))
+    xo.assert_allclose(tw.f1010, rdt_mad_at_s['f1010'], rtol=1e-3,
+                    atol=5e-3 * np.max(np.abs(rdt_mad_at_s['f1010'])))
+
 
 
 def compute_rdt(r11, r12, r21, r22, betx, bety, alfx, alfy):

{xtrack-0.98.12 → xtrack-0.98.14}/tests/test_quadrupole_fringe_ptc.py

@@ -20,7 +20,7 @@ def test_quadrupole_fringe_ptc():
 
     p0 = xt.Particles(x=x0,px=px0,y=y0,py=py0,delta=delta0,zeta=zeta0,beta0=beta0)
 
-    ptau0 = float(p0.ptau)
+    ptau0 = p0.ptau[0]
     tau0 = zeta0/beta0
 
     # XSuite
@@ -58,9 +58,9 @@ def test_quadrupole_fringe_ptc():
 
     df = madx.table.tracksumm.dframe()
 
-    assert np.isclose(p0.x, df.x[-1])
-    assert np.isclose(p0.px, df.px[-1])
-    assert np.isclose(p0.y, df.y[-1])
-    assert np.isclose(p0.py, df.py[-1])
-    assert np.isclose(p0.ptau, df.pt[-1])
-    assert np.isclose(p0.zeta/p0.beta0, df.t[-1])
+    assert np.isclose(p0.x, df.x.values[-1])
+    assert np.isclose(p0.px, df.px.values[-1])
+    assert np.isclose(p0.y, df.y.values[-1])
+    assert np.isclose(p0.py, df.py.values[-1])
+    assert np.isclose(p0.ptau, df.pt.values[-1])
+    assert np.isclose(p0.zeta/p0.beta0, df.t.values[-1])

xtrack-0.98.14/xtrack/_version.py

@@ -0,0 +1 @@
+__version__ = '0.98.14'

{xtrack-0.98.12 → xtrack-0.98.14}/xtrack/twiss.py

@@ -944,6 +944,8 @@ def twiss_line(line, particle_ref=None, method=None,
         alfy1, alfy2 = twiss_res['alfy1'], twiss_res['alfy2']
         coupling_result = _compute_coupling_elements_edwards_teng(
             W_matrix=twiss_res['W_matrix'],
+            mux=twiss_res['mux'],
+            muy=twiss_res['muy'],
             qx=twiss_res['qx'],
             qy=twiss_res['qy']
         )
@@ -1409,100 +1411,57 @@ def _compute_lattice_functions(Ws, use_full_inverse, s_co):
 
 def _compute_coupling_elements_edwards_teng(
         W_matrix: np.ndarray,
+        mux: np.ndarray,
+        muy: np.ndarray,
         qx: float = None,
         qy: float = None,
 ):
     """Compute coupling matrix elements using the Edwards-Teng method.
 
-    Using definition in chapter 7 of MAD8 guide
-
-    For the RDTs calculation, refer to R. Calaga and R. Tomás, "Betatron coupling:
-    Merging Hamiltonian and matrix approaches", 10.1103/PhysRevSTAB.8.034001.
-
-    Returns
-    -------
-    r11, r12, r21, r22: complex
-        Coupling matrix elements.
-    f1010, f1001: complex
-        Resonance driving terms f1010 and f1001.
     """
 
-    Rot = np.zeros(shape=(6, 6), dtype=np.float64)
-    lnf = xt.linear_normal_form
-
-    Rot[0:2,0:2] = lnf.Rot2D(qx)
-    Rot[2:4,2:4] = lnf.Rot2D(qy)
-
-    num_places = W_matrix.shape[0]
-    r11 = np.zeros(num_places)
-    r12 = np.zeros(num_places)
-    r21 = np.zeros(num_places)
-    r22 = np.zeros(num_places)
-    betx_et = np.zeros(num_places)
-    bety_et = np.zeros(num_places)
-    alfx_et = np.zeros(num_places)
-    alfy_et = np.zeros(num_places)
-    for idx in range(num_places):
-
-        WW = W_matrix[idx, :, :]
+    # This computes edwards-teng parameters from full one-turn W matrix at all locations
+    # edw_teng_cols = _edwards_teng_from_one_turn_at_all_locations(W_matrix, qx, qy)
+    #
+    # The following instead computes from the one-turn R matrix at one location
+    # and then propagates along the ring:
 
-        WW_inv = lnf.S.T @ WW.T @ lnf.S
-
-        RR = WW @ Rot @ WW_inv
-
-        AA = RR[:2, :2]
-        BB = RR[:2, 2:4]
-        CC = RR[2:4, :2]
-        DD = RR[2:4, 2:4]
+    # R matrix of the full ring (4D)
+    Rot = np.zeros(shape=(6, 6), dtype=np.float64)
+    Rot[0:2,0:2] = lnf.Rot2D(2 * np.pi * qx)
+    Rot[2:4,2:4] = lnf.Rot2D(2 * np.pi * qy)
+    WW0 = W_matrix[0, :, :]
+    WW0_inv = lnf.S.T @ WW0.T @ lnf.S
+    RR = WW0 @ Rot @ WW0_inv
+
+    # Edwards-Teng initial conditions
+    edw_teng_init = _compute_edwards_teng_initial(RR)
+
+    # Edwards-Teng parameters along the ring
+    edw_teng_cols = _propagate_edwards_teng(
+        WW=W_matrix, mux=mux, muy=muy,
+        RR_ET0=edw_teng_init['RR_ET0'],
+        betx0=edw_teng_init['betx0'],
+        alfx0=edw_teng_init['alfx0'],
+        bety0=edw_teng_init['bety0'],
+        alfy0=edw_teng_init['alfy0']
+    )
 
-        if np.linalg.norm(BB) < 1e-10 and np.linalg.norm(CC) < 1e-10:
-            R_edw_teng = np.zeros((2, 2))
-        else:
-            tr = np.linalg.trace
-            b_pl_c = CC + _conj_mat(BB)
-            det_bc = np.linalg.det(b_pl_c)
-            tr_a_m_tr_d = tr(AA) - tr(DD)
-            coeff = - (0.5 * tr_a_m_tr_d
-                + np.sign(det_bc) * np.sqrt(det_bc + 0.25 * tr_a_m_tr_d**2))
-            R_edw_teng = 1/coeff * b_pl_c
-
-        EE = AA - BB@R_edw_teng
-        FF = DD + R_edw_teng@BB
-
-        quarter = 0.25
-        two = 2.0
-
-        sinmu2 = -EE[0,1]*EE[1,0] - quarter*(EE[0,0] - EE[1,1])**2
-        sinmux = np.sign(EE[0,1]) * np.sqrt(abs(sinmu2))
-        betx_et_this = EE[0,1] / sinmux
-        alfx_et_this = (EE[0,0] - EE[1,1]) / (two * sinmux)
-
-        sinmu2 = -FF[0,1]*FF[1,0] - quarter*(FF[0,0] - FF[1,1])**2
-        sinmuy = np.sign(FF[0,1]) * np.sqrt(abs(sinmu2))
-        bety_et_this = FF[0,1] / sinmuy
-        alfy_et_this = (FF[0,0] - FF[1,1]) / (two * sinmuy)
-
-        r11[idx] = R_edw_teng[0,0]
-        r12[idx] = R_edw_teng[0,1]
-        r21[idx] = R_edw_teng[1,0]
-        r22[idx] = R_edw_teng[1,1]
-        betx_et[idx] = betx_et_this
-        alfx_et[idx] = alfx_et_this
-        bety_et[idx] = bety_et_this
-        alfy_et[idx] = alfy_et_this
-
-    rdts = _compute_coupling_rdts(r11, r12, r21, r22,
-                                   betx_et, bety_et, alfx_et, alfy_et)
+    # Coupling RDTs from Edwards-Teng parameters
+    rdts = _compute_coupling_rdts(edw_teng_cols['r11'], edw_teng_cols['r12'],
+                                  edw_teng_cols['r21'], edw_teng_cols['r22'],
+                                  edw_teng_cols['betx'], edw_teng_cols['bety'],
+                                  edw_teng_cols['alfx'], edw_teng_cols['alfy'])
 
     out = {
-        'r11_edw_teng': r11,
-        'r12_edw_teng': r12,
-        'r21_edw_teng': r21,
-        'r22_edw_teng': r22,
-        'betx_edw_teng': betx_et,
-        'alfx_edw_teng': alfx_et,
-        'bety_edw_teng': bety_et,
-        'alfy_edw_teng': alfy_et,
+        'r11_edw_teng': edw_teng_cols['r11'],
+        'r12_edw_teng': edw_teng_cols['r12'],
+        'r21_edw_teng': edw_teng_cols['r21'],
+        'r22_edw_teng': edw_teng_cols['r22'],
+        'betx_edw_teng': edw_teng_cols['betx'],
+        'alfx_edw_teng': edw_teng_cols['alfx'],
+        'bety_edw_teng': edw_teng_cols['bety'],
+        'alfy_edw_teng': edw_teng_cols['alfy'],
     }
     out.update(rdts)
 
@@ -1556,6 +1515,50 @@ def _compute_coupling_rdts(r11, r12, r21, r22, betx, bety, alfx, alfy):
 
     return {'f1001': f1001, 'f1010': f1010}
 
+def _compute_edwards_teng_initial(RR):
+
+    AA = RR[:2, :2]
+    BB = RR[:2, 2:4]
+    CC = RR[2:4, :2]
+    DD = RR[2:4, 2:4]
+
+    if np.linalg.norm(BB) < 1e-10 and np.linalg.norm(CC) < 1e-10:
+        RR_ET0 = np.zeros((2, 2))
+    else:
+        tr = np.linalg.trace
+        b_pl_c = CC + _conj_mat(BB)
+        det_bc = np.linalg.det(b_pl_c)
+        tr_a_m_tr_d = tr(AA) - tr(DD)
+        coeff = - (0.5 * tr_a_m_tr_d
+            + np.sign(tr_a_m_tr_d) * np.sqrt(det_bc + 0.25 * tr_a_m_tr_d**2))
+        RR_ET0 = 1/coeff * b_pl_c
+
+    EE = AA - BB@RR_ET0
+    FF = DD + RR_ET0@BB
+
+    quarter = 0.25
+    two = 2.0
+
+    sinmu2 = -EE[0,1]*EE[1,0] - quarter*(EE[0,0] - EE[1,1])**2
+    sinmux = np.sign(EE[0,1]) * np.sqrt(abs(sinmu2))
+    betx0 = EE[0,1] / sinmux
+    alfx0 = (EE[0,0] - EE[1,1]) / (two * sinmux)
+
+    sinmu2 = -FF[0,1]*FF[1,0] - quarter*(FF[0,0] - FF[1,1])**2
+    sinmuy = np.sign(FF[0,1]) * np.sqrt(abs(sinmu2))
+    bety0 = FF[0,1] / sinmuy
+    alfy0 = (FF[0,0] - FF[1,1]) / (two * sinmuy)
+
+    edw_teng_init = {
+        'RR_ET0': RR_ET0,
+        'betx0': betx0,
+        'alfx0': alfx0,
+        'bety0': bety0,
+        'alfy0': alfy0
+    }
+
+    return edw_teng_init
+
 def _conj_mat(mm):
     a = mm[0,0]
     b = mm[0,1]
@@ -1564,6 +1567,171 @@ def _conj_mat(mm):
     return np.array([[d, -b], [-c, a]])
 
 
+def _edwards_teng_from_one_turn_at_all_locations(WW, qx, qy):
+
+    # R matrix of the full ring (4D)
+    Rot = np.zeros(shape=(6, 6), dtype=np.float64)
+    Rot[0:2,0:2] = lnf.Rot2D(2 * np.pi * qx)
+    Rot[2:4,2:4] = lnf.Rot2D(2 * np.pi * qy)
+
+    n_elem = WW.shape[0]
+
+    betx = np.zeros(n_elem)
+    alfx = np.zeros(n_elem)
+    bety = np.zeros(n_elem)
+    alfy = np.zeros(n_elem)
+    r11 = np.zeros(n_elem)
+    r12 = np.zeros(n_elem)
+    r21 = np.zeros(n_elem)
+    r22 = np.zeros(n_elem)
+
+    for ii in range(n_elem):
+
+        WW0 = WW[ii, :, :]
+        WW0_inv = lnf.S.T @ WW0.T @ lnf.S
+        RR = WW0 @ Rot @ WW0_inv
+
+        # Edwards-Teng initial conditions
+        edw_teng_init = _compute_edwards_teng_initial(RR)
+
+        RR_ET=edw_teng_init['RR_ET0']
+
+        r11[ii] = RR_ET[0, 0]
+        r12[ii] = RR_ET[0, 1]
+        r21[ii] = RR_ET[1, 0]
+        r22[ii] = RR_ET[1, 1]
+        betx[ii] = edw_teng_init['betx0']
+        alfx[ii] = edw_teng_init['alfx0']
+        bety[ii] = edw_teng_init['bety0']
+        alfy[ii] = edw_teng_init['alfy0']
+
+    out_dict = {
+        'betx': betx,
+        'alfx': alfx,
+        'bety': bety,
+        'alfy': alfy,
+        'r11': r11,
+        'r12': r12,
+        'r21': r21,
+        'r22': r22
+    }
+
+    return out_dict
+
+def _propagate_edwards_teng(WW, mux, muy, RR_ET0, betx0, alfx0, bety0, alfy0):
+
+    lnf = xt.linear_normal_form
+    SS2D = lnf.S[:2, :2]
+
+    RR_ET = RR_ET0.copy()
+
+    n_elem = len(mux)
+    betx = np.zeros(n_elem)
+    alfx = np.zeros(n_elem)
+    bety = np.zeros(n_elem)
+    alfy = np.zeros(n_elem)
+    r11 = np.zeros(n_elem)
+    r12 = np.zeros(n_elem)
+    r21 = np.zeros(n_elem)
+    r22 = np.zeros(n_elem)
+
+    betx[0] = betx0
+    alfx[0] = alfx0
+    bety[0] = bety0
+    alfy[0] = alfy0
+    r11[0] = RR_ET[0, 0]
+    r12[0] = RR_ET[0, 1]
+    r21[0] = RR_ET[1, 0]
+    r22[0] = RR_ET[1, 1]
+
+    for ii in range(n_elem - 1):
+
+        # Build 2D R matrix of the element
+        WW1 = WW[ii, :, :]
+        WW2 = WW[ii+1, :, :]
+        WW1_inv = lnf.S.T @ WW1.T @ lnf.S
+        Rot_e_ii = np.zeros((6,6), dtype=np.float64)
+        Rot_e_ii[0:2,0:2] = lnf.Rot2D(2*np.pi*(mux[ii+1] - mux[ii]))
+        Rot_e_ii[2:4,2:4] = lnf.Rot2D(2*np.pi*(muy[ii+1] - muy[ii]))
+        RRe_ii = WW2 @ Rot_e_ii @ WW1_inv
+
+        # Blocks of the R matrix of the element
+        AA = RRe_ii[:2, :2]
+        BB = RRe_ii[:2, 2:4]
+        CC = RRe_ii[2:4, :2]
+        DD = RRe_ii[2:4, 2:4]
+
+        # Propagate EE, FF and RR_ET through the element
+        # Bases on MAD-X implementation (see madx/src/twiss.f90, subroutine twcptk)
+
+        if np.allclose(BB, 0, atol=1e-12) and np.allclose(CC, 0, atol=1e-12):
+            # Case in which the matrix is block diagonal (no coupling in the element)
+            EE = AA
+            FF = DD
+            EEBAR = SS2D @ EE.T @ SS2D.T
+            edet = EE[0,0]*EE[1,1] - EE[0,1]*EE[1,0]
+            CCDD = -FF @ RR_ET
+            RR_ET = -CCDD @ EEBAR / edet
+        else:
+            RR_ET_BAR = SS2D @ RR_ET.T @ SS2D.T
+            EE = AA - BB @ RR_ET
+            edet = EE[0,0]*EE[1,1] - EE[0,1]*EE[1,0]
+            EEBAR = SS2D @ EE.T @ SS2D.T
+            CCDD = CC - DD @ RR_ET
+            FF = DD + CC @ RR_ET_BAR
+            RR_ET = -CCDD @ EEBAR / edet
+
+        # Propagate Edwards-Teng Twiss parameters through the element
+        # Based on MAD-X implementation (see madx/src/twiss.f90, subroutine twcptk_twiss)
+
+        betx1 = betx[ii]
+        alfx1 = alfx[ii]
+        bety1 = bety[ii]
+        alfy1 = alfy[ii]
+
+        Rx11 = EE[0,0]
+        Rx12 = EE[0,1]
+        Rx21 = EE[1,0]
+        Rx22 = EE[1,1]
+        detx = Rx11 * Rx22 - Rx12 * Rx21
+        tempb = Rx11 * betx1 - Rx12 * alfx1
+        tempa = Rx21 * betx1 - Rx22 * alfx1
+        alfx2 = - (tempa * tempb + Rx12 * Rx22) / (detx*betx1)
+        betx2 =   (tempb * tempb + Rx12 * Rx12) / (detx*betx1)
+
+        Ry11 = FF[0,0]
+        Ry12 = FF[0,1]
+        Ry21 = FF[1,0]
+        Ry22 = FF[1,1]
+        dety = Ry11 * Ry22 - Ry12 * Ry21
+        tempb = Ry11 * bety1 - Ry12 * alfy1
+        tempa = Ry21 * bety1 - Ry22 * alfy1
+        alfy2 = - (tempa * tempb + Ry12 * Ry22) / (dety*bety1)
+        bety2 =   (tempb * tempb + Ry12 * Ry12) / (dety*bety1)
+
+        betx[ii+1] = betx2
+        alfx[ii+1] = alfx2
+        r11[ii+1] = RR_ET[0, 0]
+        r12[ii+1] = RR_ET[0, 1]
+        r21[ii+1] = RR_ET[1, 0]
+        r22[ii+1] = RR_ET[1, 1]
+        bety[ii+1] = bety2
+        alfy[ii+1] = alfy2
+
+    out_dict = {
+        'betx': betx,
+        'alfx': alfx,
+        'bety': bety,
+        'alfy': alfy,
+        'r11': r11,
+        'r12': r12,
+        'r21': r21,
+        'r22': r22
+    }
+
+    return out_dict
+
+
 def _compute_global_quantities(line, twiss_res, method):
 
         s_vect = twiss_res['s']

{xtrack-0.98.12 → xtrack-0.98.14/xtrack.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xtrack
-Version: 0.98.12
+Version: 0.98.14
 Summary: Tracking library for particle accelerators
 Author: G. Iadarola et al.
 License-Expression: Apache-2.0

xtrack-0.98.12/xtrack/_version.py

@@ -1 +0,0 @@
-__version__ = '0.98.12'