xtrack 0.92.0__tar.gz → 0.93.0__tar.gz

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

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

xtrack-0.93.0/tests/test_boris_spatial.py

@@ -0,0 +1,117 @@
+import numpy as np
+from scipy.constants import c as clight
+from scipy.constants import e as qe
+
+import xobjects as xo
+import xtrack as xt
+from xtrack._temp.boris_and_solenoid_map.solenoid_field import SolenoidField
+
+
+def test_boris_spatial():
+    delta=np.array([0, 4])
+    p0 = xt.Particles(mass0=xt.ELECTRON_MASS_EV, q0=1,
+                    energy0=45.6e9/1000,
+                    x=[-1e-3, -1e-3],
+                    px=-1e-3*(1+delta),
+                    y=1e-3,
+                    delta=delta)
+
+    p = p0.copy()
+
+    sf = SolenoidField(L=4, a=0.3, B0=1.5, z0=20)
+
+    integrator = xt.BorisSpatialIntegrator(fieldmap_callable=sf.get_field,
+                                            s_start=0,
+                                            s_end=30,
+                                            n_steps=15000)
+    p_boris = p.copy()
+    integrator.track(p_boris)
+
+
+    x_log = np.array(integrator.x_log)
+    y_log = np.array(integrator.y_log)
+    z_log = np.array(integrator.z_log)
+
+    b_seen = sf.get_field(x_log, y_log, z_log)
+
+
+    z_axis = np.linspace(0, 30, 1001)
+    Bz_axis = sf.get_field(0 * z_axis, 0 * z_axis, z_axis)[2]
+
+    P0_J = p.p0c[0] * qe / clight
+    brho = P0_J / qe / p.q0
+
+    # ks = 0.5 * (Bz_axis[:-1] + Bz_axis[1:]) / brho
+    ks = Bz_axis / brho
+    ks_entry = ks[:-1]
+    ks_exit = ks_entry*0
+    ks_exit = ks[1:]
+
+    dz = z_axis[1]-z_axis[0]
+
+    line = xt.Line(elements=[xt.VariableSolenoid(length=dz,
+                                        ks_profile=[ks_entry[ii], ks_exit[ii]])
+                                for ii in range(len(z_axis)-1)])
+    line.build_tracker()
+
+    p_xt = p0.copy()
+    line.track(p_xt, turn_by_turn_monitor='ONE_TURN_EBE')
+    mon = line.record_last_track
+
+    p_xt = p0.copy()
+    line.configure_radiation(model=None)
+    line.track(p_xt, turn_by_turn_monitor='ONE_TURN_EBE')
+    mon_no_rad = line.record_last_track
+
+    Bz_mid = 0.5 * (Bz_axis[:-1] + Bz_axis[1:])
+    Bz_mon = 0 * Bz_axis
+    Bz_mon[1:] = Bz_mid
+
+    # Wolsky Eq. 3.114
+    Ax = -0.5 * Bz_mon * mon.y
+    Ay =  0.5 * Bz_mon * mon.x
+
+    # Wolsky Eq. 2.74
+    ax_ref = Ax * p0.q0 * qe / P0_J
+    ay_ref = Ay * p0.q0 * qe / P0_J
+
+    px_mech = mon.px - ax_ref
+    py_mech = mon.py - ay_ref
+    pz_mech = np.sqrt((1 + mon.delta)**2 - px_mech**2 - py_mech**2)
+
+    xp = px_mech / pz_mech
+    yp = py_mech / pz_mech
+
+    dx_ds = np.diff(mon.x, axis=1) / np.diff(mon.s, axis=1)
+    dy_ds = np.diff(mon.y, axis=1) / np.diff(mon.s, axis=1)
+
+    z_check = sf.z0 + sf.L * np.linspace(-2, 2, 1001)
+
+    for i_part in range(z_log.shape[1]):
+
+        this_s_boris = 0.5 * (z_log[:-1, i_part] + z_log[1:, i_part])
+        dx_ds_boris = np.diff(x_log[:, i_part]) / np.diff(z_log[:, i_part])
+        dy_ds_boris = np.diff(y_log[:, i_part]) / np.diff(z_log[:, i_part])
+
+        s_xsuite = 0.5 * (mon.s[i_part, :-1] + mon.s[i_part, 1:])
+        dx_ds_xsuite = np.diff(mon.x[i_part, :]) / np.diff(mon.s[i_part, :])
+        dy_ds_xsuite = np.diff(mon.y[i_part, :]) / np.diff(mon.s[i_part, :])
+
+        dx_ds_xsuite_check = np.interp(z_check, s_xsuite, dx_ds_xsuite)
+        dy_ds_xsuite_check = np.interp(z_check, s_xsuite, dy_ds_xsuite)
+
+        dx_ds_boris_check = np.interp(z_check, this_s_boris, dx_ds_boris)
+        dy_ds_boris_check = np.interp(z_check, this_s_boris, dy_ds_boris)
+
+        this_dx_ds = dx_ds[i_part, :]
+        this_dy_ds = dy_ds[i_part, :]
+
+        xo.assert_allclose(dx_ds_xsuite_check, dx_ds_boris_check, rtol=0,
+                atol=2.8e-2 * (np.max(dx_ds_boris_check) - np.min(dx_ds_boris_check)))
+        xo.assert_allclose(dy_ds_xsuite_check, dy_ds_boris_check, rtol=0,
+                atol=2.8e-2 * (np.max(dy_ds_boris_check) - np.min(dy_ds_boris_check)))
+
+        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))

{xtrack-0.92.0 → xtrack-0.93.0}/xtrack/__init__.py

@@ -47,6 +47,7 @@ from .trajectory_correction import TrajectoryCorrection
 from .mad_parser.loader import load_madx_lattice
 from . import json
 from .load import load
+from .boris import BorisSpatialIntegrator
 
 from .multiline import Multiline
 

xtrack-0.93.0/xtrack/_version.py

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

xtrack-0.93.0/xtrack/boris.py

@@ -0,0 +1,161 @@
+import numpy as np
+
+from scipy.constants import c as clight
+from scipy.constants import e as qe
+
+class BorisSpatialIntegrator:
+
+    isthick = True
+
+    def __init__(self, fieldmap_callable, s_start, s_end, n_steps, verbose=False):
+        self.fieldmap_callable = fieldmap_callable
+        self.s_start = s_start
+        self.s_end = s_end
+        self.ds = (s_end - s_start) / n_steps
+        self.n_steps = n_steps
+        self.verbose = verbose
+        self.length = s_end - s_start
+
+    def track(self, p):
+
+        x_log = []
+        y_log = []
+        z_log = []
+        s_in = p.s.copy()
+        p.s=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
+
+            P0 = p0c * qe / clight # in kg m/s
+            P = P0 * (1 + delta)
+
+            Px = px * P0
+            Py = py * P0
+
+            w = np.zeros((Px.shape[0], 3), dtype=Px.dtype)
+            w[:, 0] = Px
+            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_log.append(p.x.copy())
+            y_log.append(p.y.copy())
+            z_log.append(p.s.copy())
+        p.s = 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)
+
+
+import numpy as np
+c = 299_792_458.0  # m/s
+
+def step_spatial_boris_B(x, y, z, w, q, dz, field_fn):
+    """
+    Spatial Boris step for magnetic fields only (E = 0),
+    using constant total momentum magnitude P = w[:,2].
+
+    Parameters
+    ----------
+    x, y, z : (N,) arrays
+        Particle positions [m].
+    w : (N, 3) array
+        Momentum-like state: (px, py, P), where
+        P = |p| = sqrt(px^2 + py^2 + pz^2) is constant.
+    q : float
+        Particle charge [C].
+    dz : float
+        Spatial step along z [m].
+    field_fn : callable(x, y, z) -> (Bx, By, Bz)
+        Function returning magnetic-field components [T].
+
+    Returns
+    -------
+    x1, y1, z1 : (N,) arrays
+        Updated positions [m].
+    w1 : (N, 3) array
+        Updated momentum state (px, py, P).
+    """
+
+    # --- Unpack and ensure arrays
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+    w = np.asarray(w)
+    px, py, P = w[:, 0], w[:, 1], w[:, 2]
+
+    # --- Compute longitudinal momentum from constant magnitude
+    pz = np.sqrt(np.maximum(P**2 - px**2 - py**2, 0.0))
+
+    # --- Half drift in (x, y)
+    xh = x + (px / pz) * (dz * 0.5)
+    yh = y + (py / pz) * (dz * 0.5)
+    zh = z + dz * 0.5
+
+    # --- Evaluate magnetic field at mid-step
+    Bx, By, Bz = field_fn(xh, yh, zh)
+
+    # ============================================
+    # (1) FIRST HALF-KICK from (Bx, By)
+    # ============================================
+    pxm = px - 0.5 * q * dz * By
+    pym = py + 0.5 * q * dz * Bx
+
+    # --- Recompute pz to maintain |p| = P
+    pz = np.sqrt(np.maximum(P**2 - pxm**2 - pym**2, 0.0))
+
+    # ============================================
+    # (2) ROTATION due to Bz
+    # ============================================
+    t = 0.5 * q * Bz * dz / pz
+    t2 = t * t
+    s = 2.0 * t / (1.0 + t2)
+    c0 = (1.0 - t2) / (1.0 + t2)
+
+    # pxp = c0 * pxm - s * pym
+    # pyp = s * pxm + c0 * pym
+    pxp = c0 * pxm + s * pym
+    pyp = -s * pxm + c0 * pym
+
+    # ============================================
+    # (3) SECOND HALF-KICK from (Bx, By)
+    # ============================================
+    px1 = pxp - 0.5 * q * dz * By
+    py1 = pyp + 0.5 * q * dz * Bx
+
+    # --- Recompute pz after full step
+    pz1 = np.sqrt(np.maximum(P**2 - px1**2 - py1**2, 0.0))
+
+    # ============================================
+    # (4) SECOND HALF-DRIFT (x, y)
+    # ============================================
+    x1 = xh + (px1 / pz1) * (dz * 0.5)
+    y1 = yh + (py1 / pz1) * (dz * 0.5)
+    z1 = z + dz
+
+    # --- Pack updated (px, py, P)
+    w1 = np.stack([px1, py1, P], axis=1)
+
+    return x1, y1, z1, w1

{xtrack-0.92.0 → xtrack-0.93.0}/xtrack/line.py

@@ -1044,6 +1044,8 @@ class Line:
 
     @particle_ref.setter
     def particle_ref(self, particle_ref):
+        if isinstance(particle_ref, LineParticleRef):
+            particle_ref = particle_ref.line._particle_ref
         self._particle_ref = particle_ref
         # This looks a bit dangerous, when working with coasting beams in environments.
         # If the particle is shared with other lines, t_sim might be wrong.
@@ -2514,7 +2516,7 @@ class Line:
         s_cuts = list(tab_insertions['s_start']) + list(tab_insertions['s_end'])
         s_cuts = list(set(s_cuts))
 
-        self.cut_at_s(s_cuts, s_tol=1e-06, return_slices=True)
+        self.cut_at_s(s_cuts, s_tol=s_tol, return_slices=True)
 
         tt_after_cut = self.get_table()
         tt_after_cut['length'] = np.diff(tt_after_cut.s, append=tt_after_cut.s[-1])

{xtrack-0.92.0 → xtrack-0.93.0}/xtrack/twiss.py

@@ -1595,105 +1595,108 @@ def _compute_chromatic_functions(line, init, delta_chrom, steps_r_matrix,
                     only_markers=False,
                     periodic=False,
                     periodic_mode=None,
-                    include_collective=False):
+                    include_collective=False,
+                    tw_chrom_res=None
+                    ):
 
     if only_markers:
         raise NotImplementedError('only_markers not supported anymore')
 
-    tw_chrom_res = []
-    for dd in [-delta_chrom, delta_chrom]:
-        tw_init_chrom = init.copy()
-
-        if periodic:
-            import xpart
-            part_guess = xpart.build_particles(
-                _context=line._context,
-                x_norm=0,
-                zeta=tw_init_chrom.zeta,
-                delta=tw_init_chrom.delta+ dd,
-                particle_on_co=on_momentum_twiss_res.particle_on_co.copy(),
-                nemitt_x=nemitt_x, nemitt_y=nemitt_y,
-                W_matrix=tw_init_chrom.W_matrix,
-                include_collective=include_collective)
-            part_chrom = line.find_closed_orbit(delta0=dd, co_guess=part_guess,
-                                    start=start, end=end, num_turns=num_turns,
-                                    symmetrize=False,
-                                    include_collective=include_collective,
-                                    )
-            tw_init_chrom.particle_on_co = part_chrom
-            RR_chrom = line.compute_one_turn_matrix_finite_differences(
-                                        particle_on_co=tw_init_chrom.particle_on_co.copy(),
+    if tw_chrom_res is None:
+        tw_chrom_res = []
+        for dd in [-delta_chrom, delta_chrom]:
+            tw_init_chrom = init.copy()
+
+            if periodic:
+                import xpart
+                part_guess = xpart.build_particles(
+                    _context=line._context,
+                    x_norm=0,
+                    zeta=tw_init_chrom.zeta,
+                    delta=tw_init_chrom.delta+ dd,
+                    particle_on_co=on_momentum_twiss_res.particle_on_co.copy(),
+                    nemitt_x=nemitt_x, nemitt_y=nemitt_y,
+                    W_matrix=tw_init_chrom.W_matrix,
+                    include_collective=include_collective)
+                part_chrom = line.find_closed_orbit(delta0=dd, co_guess=part_guess,
                                         start=start, end=end, num_turns=num_turns,
-                                        steps_r_matrix=steps_r_matrix,
                                         symmetrize=False,
                                         include_collective=include_collective,
-                                        )['R_matrix']
-            (WW_chrom, _, _, _) = lnf.compute_linear_normal_form(RR_chrom,
-                                    only_4d_block=True,
-                                    responsiveness_tol=matrix_responsiveness_tol,
-                                    stability_tol=matrix_stability_tol,
-                                    symplectify=symplectify)
-            tw_init_chrom.W_matrix = WW_chrom
-        else:
-            alfx = init.alfx
-            betx = init.betx
-            alfy = init.alfy
-            bety = init.bety
-            dx = init.dx
-            dy = init.dy
-            dpx = init.dpx
-            dpy = init.dpy
-            ddx = init.ddx
-            ddpx = init.ddpx
-            ddy = init.ddy
-            ddpy = init.ddpy
-            ax_chrom = init.ax_chrom
-            bx_chrom = init.bx_chrom
-            ay_chrom = init.ay_chrom
-            by_chrom = init.by_chrom
-
-            dbetx_dpzeta = bx_chrom * betx
-            dbety_dpzeta = by_chrom * bety
-            dalfx_dpzeta = ax_chrom + bx_chrom * alfx
-            dalfy_dpzeta = ay_chrom + by_chrom * alfy
-
-            tw_init_chrom.particle_on_co.x += dx * dd + 1/2 * ddx * dd**2
-            tw_init_chrom.particle_on_co.px += dpx * dd + 1/2 * ddpx * dd**2
-            tw_init_chrom.particle_on_co.y += dy * dd + 1/2 * ddy * dd**2
-            tw_init_chrom.particle_on_co.py += dpy * dd + 1/2 * ddpy * dd**2
-            tw_init_chrom.particle_on_co.delta += dd
-
-            twinit_aux = TwissInit(
-                alfx=alfx + dalfx_dpzeta * dd,
-                betx=betx + dbetx_dpzeta * dd,
-                alfy=alfy + dalfy_dpzeta * dd,
-                bety=bety + dbety_dpzeta * dd,
-                dx=dx + ddx * dd,
-                dpx=dpx + ddpx * dd,
-                dy=dy + ddy * dd,
-                dpy=dpy + ddpy * dd)
-            twinit_aux._complete(line, element_name=init.element_name)
-            tw_init_chrom.W_matrix = twinit_aux.W_matrix
-
-        tw_chrom_res.append(
-            _twiss_open(
-                line=line,
-                init=tw_init_chrom,
-                start=start, end=end,
-                nemitt_x=nemitt_x,
-                nemitt_y=nemitt_y,
-                r_sigma=r_sigma,
-                delta_disp=delta_disp,
-                use_full_inverse=use_full_inverse,
-                hide_thin_groups=hide_thin_groups,
-                only_markers=only_markers,
-                _continue_if_lost=False,
-                _keep_tracking_data=False,
-                _keep_initial_particles=False,
-                _initial_particles=None,
-                _ebe_monitor=None,
+                                        )
+                tw_init_chrom.particle_on_co = part_chrom
+                RR_chrom = line.compute_one_turn_matrix_finite_differences(
+                                            particle_on_co=tw_init_chrom.particle_on_co.copy(),
+                                            start=start, end=end, num_turns=num_turns,
+                                            steps_r_matrix=steps_r_matrix,
+                                            symmetrize=False,
+                                            include_collective=include_collective,
+                                            )['R_matrix']
+                (WW_chrom, _, _, _) = lnf.compute_linear_normal_form(RR_chrom,
+                                        only_4d_block=True,
+                                        responsiveness_tol=matrix_responsiveness_tol,
+                                        stability_tol=matrix_stability_tol,
+                                        symplectify=symplectify)
+                tw_init_chrom.W_matrix = WW_chrom
+            else:
+                alfx = init.alfx
+                betx = init.betx
+                alfy = init.alfy
+                bety = init.bety
+                dx = init.dx
+                dy = init.dy
+                dpx = init.dpx
+                dpy = init.dpy
+                ddx = init.ddx
+                ddpx = init.ddpx
+                ddy = init.ddy
+                ddpy = init.ddpy
+                ax_chrom = init.ax_chrom
+                bx_chrom = init.bx_chrom
+                ay_chrom = init.ay_chrom
+                by_chrom = init.by_chrom
+
+                dbetx_dpzeta = bx_chrom * betx
+                dbety_dpzeta = by_chrom * bety
+                dalfx_dpzeta = ax_chrom + bx_chrom * alfx
+                dalfy_dpzeta = ay_chrom + by_chrom * alfy
+
+                tw_init_chrom.particle_on_co.x += dx * dd + 1/2 * ddx * dd**2
+                tw_init_chrom.particle_on_co.px += dpx * dd + 1/2 * ddpx * dd**2
+                tw_init_chrom.particle_on_co.y += dy * dd + 1/2 * ddy * dd**2
+                tw_init_chrom.particle_on_co.py += dpy * dd + 1/2 * ddpy * dd**2
+                tw_init_chrom.particle_on_co.delta += dd
+
+                twinit_aux = TwissInit(
+                    alfx=alfx + dalfx_dpzeta * dd,
+                    betx=betx + dbetx_dpzeta * dd,
+                    alfy=alfy + dalfy_dpzeta * dd,
+                    bety=bety + dbety_dpzeta * dd,
+                    dx=dx + ddx * dd,
+                    dpx=dpx + ddpx * dd,
+                    dy=dy + ddy * dd,
+                    dpy=dpy + ddpy * dd)
+                twinit_aux._complete(line, element_name=init.element_name)
+                tw_init_chrom.W_matrix = twinit_aux.W_matrix
+
+            tw_chrom_res.append(
+                _twiss_open(
+                    line=line,
+                    init=tw_init_chrom,
+                    start=start, end=end,
+                    nemitt_x=nemitt_x,
+                    nemitt_y=nemitt_y,
+                    r_sigma=r_sigma,
+                    delta_disp=delta_disp,
+                    use_full_inverse=use_full_inverse,
+                    hide_thin_groups=hide_thin_groups,
+                    only_markers=only_markers,
+                    _continue_if_lost=False,
+                    _keep_tracking_data=False,
+                    _keep_initial_particles=False,
+                    _initial_particles=None,
+                    _ebe_monitor=None,
+                )
             )
-        )
 
     dmux = (tw_chrom_res[1].mux - tw_chrom_res[0].mux)/(2*delta_chrom)
     dmuy = (tw_chrom_res[1].muy - tw_chrom_res[0].muy)/(2*delta_chrom)

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

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

{xtrack-0.92.0 → xtrack-0.93.0}/xtrack.egg-info/SOURCES.txt

@@ -26,6 +26,7 @@ tests/test_aperture_table.py
 tests/test_aperture_turn_ele_and_monitor.py
 tests/test_apertures.py
 tests/test_attr_replicas_and_slices.py
+tests/test_boris_spatial.py
 tests/test_bucket_with_ref_energy_change.py
 tests/test_cavity_absolute_time.py
 tests/test_chromatic_functions_vs_madx.py
@@ -125,6 +126,7 @@ xtrack/__init__.py
 xtrack/_version.py
 xtrack/aperture_meas.py
 xtrack/base_element.py
+xtrack/boris.py
 xtrack/environment.py
 xtrack/footprint.py
 xtrack/functions.py

xtrack-0.92.0/xtrack/_version.py

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