xtrack 0.93.0__tar.gz → 0.93.1__tar.gz

{xtrack-0.93.0/xtrack.egg-info → xtrack-0.93.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.93.0
+Version: 0.93.1
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack

{xtrack-0.93.0 → xtrack-0.93.1}/tests/test_boris_spatial.py

@@ -7,7 +7,7 @@ import xtrack as xt
 from xtrack._temp.boris_and_solenoid_map.solenoid_field import SolenoidField
 
 
-def test_boris_spatial():
+def test_boris_spatial_solenoid_map():
     delta=np.array([0, 4])
     p0 = xt.Particles(mass0=xt.ELECTRON_MASS_EV, q0=1,
                     energy0=45.6e9/1000,
@@ -114,4 +114,49 @@ def test_boris_spatial():
         xo.assert_allclose(ax_ref[i_part, :], mon.ax[i_part, :],
                         rtol=0, atol=np.max(np.abs(ax_ref)*3e-2))
         xo.assert_allclose(ay_ref[i_part, :], mon.ay[i_part, :],
-                        rtol=0, atol=np.max(np.abs(ay_ref)*3e-2))
+                        rtol=0, atol=np.max(np.abs(ay_ref)*3e-2))
+
+def test_boris_spatial_drift():
+    def zero_field(x, y, z):
+        return (0*x, 0*y, 0*z)
+
+    integrator = xt.BorisSpatialIntegrator(fieldmap_callable=zero_field,
+                                            s_start=10,
+                                            s_end=20,
+                                            n_steps=15000)
+
+    drift = xt.Drift(length=10, model='exact')
+
+    p0 = xt.Particles('proton',
+            x=[1e-2, -1e-2], px=[15e-3, 0], delta=[0, 2], p0c=1e9, zeta=1)
+
+    p_integ = p0.copy()
+    integrator.track(p_integ)
+
+    p_drift = p0.copy()
+    drift.track(p_drift)
+
+    xo.assert_allclose(p_integ.x, p_drift.x, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.px, p_drift.px, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.y, p_drift.y, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.py, p_drift.py, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.s, p_drift.s, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.delta, p_drift.delta, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.zeta, p_drift.zeta, rtol=0, atol=1e-10)
+
+    # Check behavior with lost particles
+    p0 = xt.Particles('proton',
+            x=[1e-2, -1e-2], px=[15e-3, 0], delta=[0, 2], p0c=1e9, zeta=1,
+            state=[1, -1])
+
+    p_integ = p0.copy()
+    integrator.track(p_integ)
+    p_drift = p0.copy()
+    drift.track(p_drift)
+
+    xo.assert_allclose(p_integ.x, p_drift.x, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.px, p_drift.px, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.y, p_drift.y, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.py, p_drift.py, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.s, p_drift.s, rtol=0, atol=1e-10)
+    xo.assert_allclose(p_integ.delta, p_drift.delta, rtol=0, atol=1e-10)

{xtrack-0.93.0 → xtrack-0.93.1}/tests/test_madnginterface.py

@@ -163,4 +163,142 @@ def test_madng_conversion_drift_slice():
     tw_ng = line.madng_twiss(normal_form=False)
 
     xo.assert_allclose(tw_ng.beta11_ng, tw.betx, rtol=1e-8)
-    xo.assert_allclose(tw_ng.beta22_ng, tw.bety, rtol=1e-8)
+    xo.assert_allclose(tw_ng.beta22_ng, tw.bety, rtol=1e-8)
+
+def test_madng_interface_with_slicing():
+    line = xt.load(test_data_folder /
+                            'hllhc15_thick/lhc_thick_with_knobs.json')
+
+    # Cut line at every 1 up to s = 1000
+    line.cut_at_s(np.arange(1000))
+
+    tw_xs = line.twiss4d()
+    tw = line.madng_twiss()
+
+    assert len(tw) == len(tw_xs)
+
+    xo.assert_allclose(tw.x, 0, atol=1e-10, rtol=0)
+    xo.assert_allclose(tw.y, 0, atol=1e-10, rtol=0)
+    xo.assert_allclose(tw.betx2, 0, atol=1e-10, rtol=0)
+    xo.assert_allclose(tw.bety1, 0, atol=1e-10, rtol=0)
+    xo.assert_allclose(tw.x, tw.x_ng, atol=1e-8, rtol=0)
+    xo.assert_allclose(tw.y, tw.y_ng, atol=1e-10, rtol=0)
+    xo.assert_allclose(tw.betx2, tw.beta12_ng, atol=1e-10, rtol=0)
+    xo.assert_allclose(tw.bety1, tw.beta21_ng, atol=1e-19, rtol=0)
+    xo.assert_allclose(tw.wx_chrom, tw.wx_ng, atol=5e-3*tw.wx_chrom.max(), rtol=0)
+    xo.assert_allclose(tw.wy_chrom, tw.wy_ng, atol=5e-3*tw.wy_chrom.max(), rtol=0)
+    xo.assert_allclose(tw.ax_chrom, tw.ax_ng, atol=5e-3*tw.wx_chrom.max(), rtol=0)
+    xo.assert_allclose(tw.ay_chrom, tw.ay_ng, atol=5e-3*tw.wy_chrom.max(), rtol=0)
+    xo.assert_allclose(tw.bx_chrom, tw.bx_ng, atol=5e-3*tw.wx_chrom.max(), rtol=0)
+    xo.assert_allclose(tw.by_chrom, tw.by_ng, atol=5e-3*tw.wy_chrom.max(), rtol=0)
+
+def test_madng_twiss_with_initial_conditions():
+    line = xt.load(test_data_folder /
+                            'hllhc15_thick/lhc_thick_with_knobs.json')
+    #pytest.set_trace()
+    tw_xs = line.twiss(betx=120, bety=150)
+    tw = line.madng_twiss(beta11=120, beta22=150)
+
+    assert len(tw) == len(tw_xs)
+    assert len(tw.betx) == len(tw.beta11_ng)
+
+    xo.assert_allclose(tw.betx, tw.beta11_ng, rtol=1e-7, atol=1e-6)
+    xo.assert_allclose(tw.bety, tw.beta22_ng, rtol=1e-7, atol=1e-6)
+    xo.assert_allclose(tw.alfx, tw.alfa11_ng, rtol=1e-7, atol=1e-6)
+    xo.assert_allclose(tw.alfy, tw.alfa22_ng, rtol=1e-7, atol=1e-6)
+    xo.assert_allclose(tw.dx, tw.dx_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw.dy, tw.dy_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw.dpx, tw.dpx_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw.dpy, tw.dpy_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw.x, tw.x_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw.y, tw.y_ng, rtol=1e-8, atol=1e-6)
+
+    tw2_xs = line.twiss(start='s.ds.l8.b1', end='ip1', betx=100, bety=34)
+    tw2_xsng = line.madng_twiss(start='s.ds.l8.b1', end='ip1', beta11=100, beta22=34, xsuite_tw=False)
+
+    assert len(tw2_xs.betx) == len(tw2_xsng.beta11_ng)
+    xo.assert_allclose(tw2_xs.betx, tw2_xsng.beta11_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.bety, tw2_xsng.beta22_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.alfx, tw2_xsng.alfa11_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.alfy, tw2_xsng.alfa22_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.dx, tw2_xsng.dx_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.dy, tw2_xsng.dy_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.dpx, tw2_xsng.dpx_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.dpy, tw2_xsng.dpy_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw2_xs.x, tw2_xsng.x_ng, rtol=1e-8, atol=1e-8)
+    xo.assert_allclose(tw2_xs.y, tw2_xsng.y_ng, rtol=1e-8, atol=1e-8)
+
+    tw3_xs = line.twiss(start='ip8', end='ip2', betx=1.5, bety=1.5)
+    tw3_xsng = line.madng_twiss(start='ip8', end='ip2', beta11=1.5, beta22=1.5, xsuite_tw=True)
+
+    assert len(tw3_xs.betx) == len(tw3_xsng.beta11_ng)
+    xo.assert_allclose(tw3_xs.betx, tw3_xsng.beta11_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.bety, tw3_xsng.beta22_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.alfx, tw3_xsng.alfa11_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.alfy, tw3_xsng.alfa22_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.dx, tw3_xsng.dx_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.dy, tw3_xsng.dy_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.dpx, tw3_xsng.dpx_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.dpy, tw3_xsng.dpy_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xs.x, tw3_xsng.x_ng, rtol=1e-8, atol=1e-8)
+    xo.assert_allclose(tw3_xs.y, tw3_xsng.y_ng, rtol=1e-8, atol=1e-8)
+
+    xo.assert_allclose(tw3_xsng.betx, tw3_xsng.beta11_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xsng.bety, tw3_xsng.beta22_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xsng.alfx, tw3_xsng.alfa11_ng, rtol=1e-8, atol=1e-6)
+    xo.assert_allclose(tw3_xsng.alfy, tw3_xsng.alfa22_ng, rtol=1e-8, atol=1e-6)
+
+def test_madng_slices():
+    line = xt.load(test_data_folder /
+                            'hllhc15_thick/lhc_thick_with_knobs.json')
+    tw = line.twiss4d()
+
+    twng = line.madng_twiss()
+
+    line.cut_at_s(np.linspace(0, line.get_length(), 5000))
+    tw_sliced = line.twiss4d()
+    twng_sliced = line.madng_twiss()
+    tt_sliced = line.get_table()
+
+    assert np.all(np.array(sorted(list(set(tt_sliced.element_type)))) ==
+        ['',
+        'Cavity',
+        'Drift',
+        'DriftSlice',
+        'Marker',
+        'Multipole',
+        'Octupole',
+        'Quadrupole',
+        'RBend',
+        'Sextupole',
+        'ThickSliceBend',
+        'ThickSliceCavity',
+        'ThickSliceMultipole',
+        'ThickSliceOctupole',
+        'ThickSliceQuadrupole',
+        'ThickSliceRBend',
+        'ThickSliceSextupole',
+        'ThickSliceUniformSolenoid',
+        'ThinSliceBendEntry',
+        'ThinSliceBendExit',
+        'ThinSliceOctupoleEntry',
+        'ThinSliceOctupoleExit',
+        'ThinSliceQuadrupoleEntry',
+        'ThinSliceQuadrupoleExit',
+        'ThinSliceRBendEntry',
+        'ThinSliceRBendExit',
+        'ThinSliceSextupoleEntry',
+        'ThinSliceSextupoleExit',
+        'ThinSliceUniformSolenoidEntry',
+        'ThinSliceUniformSolenoidExit',
+        'UniformSolenoid'])
+
+    twng_ip = twng.rows['ip.*']
+    twng_ip_sliced = twng_sliced.rows['ip.*']
+    xo.assert_allclose(twng_ip.s, twng_ip_sliced.s, rtol=1e-8)
+    xo.assert_allclose(twng_ip.beta11_ng, twng_ip_sliced.beta11_ng, rtol=1e-3)
+    xo.assert_allclose(twng_ip.beta22_ng, twng_ip_sliced.beta22_ng, rtol=1e-3)
+    xo.assert_allclose(twng_ip.wx_ng, twng_ip_sliced.wx_ng, rtol=1e-3)
+    xo.assert_allclose(twng_ip.wy_ng, twng_ip_sliced.wy_ng, rtol=1e-3)
+    xo.assert_allclose(twng_ip.dx_ng, twng_ip_sliced.dx_ng, atol=1e-6)
+    xo.assert_allclose(twng_ip.dy_ng, twng_ip_sliced.dy_ng, atol=1e-6)

xtrack-0.93.1/xtrack/_version.py

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

{xtrack-0.93.0 → xtrack-0.93.1}/xtrack/boris.py

@@ -18,27 +18,34 @@ class BorisSpatialIntegrator:
 
     def track(self, p):
 
+        mask_alive = p.state > 0
+
         x_log = []
         y_log = []
         z_log = []
-        s_in = p.s.copy()
-        p.s=self.s_start
+        s_in = p.s[mask_alive].copy()
+        p.s[mask_alive] = self.s_start
 
         for ii in range(self.n_steps):
 
             if self.verbose:
                 print(f's_in = {s_in[0]:.3f} s_in_map = {p.s[0]:.3f}', end='\r', flush=True)
 
-            q0 = p.q0
-            x = p.x.copy()
-            y = p.y.copy()
-            z = p.s.copy()
-            px = p.px.copy()
-            py = p.py.copy()
-            delta = p.delta.copy()
-            p0c = p.p0c
-
-            charge0_coulomb = q0 * qe
+            x = p.x[mask_alive].copy()
+            y = p.y[mask_alive].copy()
+            z = p.s[mask_alive].copy()
+            px = p.px[mask_alive].copy()
+            py = p.py[mask_alive].copy()
+            delta = p.delta[mask_alive].copy()
+            energy = p.energy[mask_alive].copy()
+            p0c = p.p0c[mask_alive].copy()
+            beta0 = p.beta0[mask_alive].copy()
+            charge = p.charge[mask_alive].copy()
+            mass = p.mass[mask_alive].copy()
+
+            charge_coulomb = charge * qe
+            mass_kg = mass * qe / clight**2
+            gamma = energy / mass
 
             P0 = p0c * qe / clight # in kg m/s
             P = P0 * (1 + delta)
@@ -51,19 +58,20 @@ class BorisSpatialIntegrator:
             w[:, 1] = Py
             w[:, 2] = P
 
-            x_new, y_new, z_new, w_new = step_spatial_boris_B(x, y, z, w,
-                charge0_coulomb, self.ds,
-                field_fn=self.fieldmap_callable)
-            p.x = x_new.copy()
-            p.y = y_new.copy()
-            p.s = z_new.copy()
-            p.px = w_new[:, 0] / P0
-            p.py = w_new[:, 1] / P0
+            x_new, y_new, z_new, w_new, dt = step_spatial_boris_B(x, y, z, w,
+                charge_coulomb, self.ds,
+                field_fn=self.fieldmap_callable, m_kg=mass_kg, gamma=gamma)
+            p.x[mask_alive] = x_new
+            p.y[mask_alive] = y_new
+            p.s[mask_alive] = z_new
+            p.px[mask_alive] = w_new[:, 0] / P0
+            p.py[mask_alive] = w_new[:, 1] / P0
+            p.zeta[mask_alive] += (self.ds - dt * clight * beta0)
 
             x_log.append(p.x.copy())
             y_log.append(p.y.copy())
             z_log.append(p.s.copy())
-        p.s = s_in + self.length
+        p.s[mask_alive] = s_in + self.length
         self.x_log = np.array(x_log)
         self.y_log = np.array(y_log)
         self.z_log = np.array(z_log)
@@ -72,7 +80,7 @@ class BorisSpatialIntegrator:
 import numpy as np
 c = 299_792_458.0  # m/s
 
-def step_spatial_boris_B(x, y, z, w, q, dz, field_fn):
+def step_spatial_boris_B(x, y, z, w, q, dz, field_fn, m_kg, gamma):
     """
     Spatial Boris step for magnetic fields only (E = 0),
     using constant total momentum magnitude P = w[:,2].
@@ -99,6 +107,8 @@ def step_spatial_boris_B(x, y, z, w, q, dz, field_fn):
         Updated momentum state (px, py, P).
     """
 
+    dt = 0.
+
     # --- Unpack and ensure arrays
     x = np.asarray(x)
     y = np.asarray(y)
@@ -113,6 +123,7 @@ def step_spatial_boris_B(x, y, z, w, q, dz, field_fn):
     xh = x + (px / pz) * (dz * 0.5)
     yh = y + (py / pz) * (dz * 0.5)
     zh = z + dz * 0.5
+    dt += (dz * 0.5) / pz * gamma * m_kg #  s
 
     # --- Evaluate magnetic field at mid-step
     Bx, By, Bz = field_fn(xh, yh, zh)
@@ -154,8 +165,9 @@ def step_spatial_boris_B(x, y, z, w, q, dz, field_fn):
     x1 = xh + (px1 / pz1) * (dz * 0.5)
     y1 = yh + (py1 / pz1) * (dz * 0.5)
     z1 = z + dz
+    dt += (dz * 0.5) / pz1 * gamma * m_kg #  s
 
     # --- Pack updated (px, py, P)
     w1 = np.stack([px1, py1, P], axis=1)
 
-    return x1, y1, z1, w1
+    return x1, y1, z1, w1, dt