Feature: Implement NonlinearKicker element

tests/test_nonlinearkicker.py

@@ -0,0 +1,105 @@
+import numpy as np
+import pytest
+
+import xobjects as xo
+import xtrack as xt
+import xpart as xp
+from xobjects.test_helpers import for_all_test_contexts
+
+
+# This decorator ensures the test runs across ContextCpu, ContextCuda, and ContextPyopencl automatically
+@for_all_test_contexts
+def test_nonlinear_kicker_vs_multiple_wires(test_context):
+
+    # ----------------------------------------------------
+    # 1. Physical Parameter Setup
+    # ----------------------------------------------------
+    p0c = 7000e9  # 7 TeV
+    l_phy = 1.5
+    l_int = 1.5
+
+    # Define 3 wires with different currents and positions
+    # Wire 1: High current, positive offset
+    # Wire 2: Negative current, negative offset
+    # Wire 3: Low current, near center
+    currents = [250.0, -150.0, 50.0]
+    x_pos = [-8e-3, 10e-3, 2.1e-3]
+    y_pos = [-10e-3, 5e-3, 4e-3]
+
+    # ----------------------------------------------------
+    # 2. Case A: Using the new NonlinearKicker element
+    # ----------------------------------------------------
+    # All wires are computed within this single element
+    kicker = xt.NonlinearKicker(
+        _context=test_context,  # Crucial: ensures allocation on the correct device memory
+        L_phy=l_phy,
+        L_int=l_int,
+        current=currents,
+        xma=x_pos,
+        yma=y_pos,
+    )
+
+    # ----------------------------------------------------
+    # 3. Case B: Using standard xt.Wire (Ground Truth Reference)
+    # ----------------------------------------------------
+    # Create 3 independent Wire elements to simulate the same effect
+    wire_elements = []
+    for i in range(len(currents)):
+        w = xt.Wire(
+            _context=test_context,
+            L_phy=l_phy,
+            L_int=l_int,
+            current=currents[i],
+            xma=x_pos[i],
+            yma=y_pos[i],
+        )
+        wire_elements.append(w)
+
+    # ----------------------------------------------------
+    # 4. Construct Test Particles
+    # ----------------------------------------------------
+    # Create a distribution of particles covering various transverse positions
+    x_test = np.array([0.0, 1e-3, -5e-3, 2e-3])
+    y_test = np.array([0.0, -2e-3, 4e-3, 1e-3])
+
+    # Particle Set 1: For testing NonlinearKicker
+    p_kicker = xp.Particles(
+        _context=test_context, p0c=p0c, x=x_test.copy(), y=y_test.copy()
+    )
+
+    # Particle Set 2: For testing Wire sequence (Reference)
+    # Must be a deep copy to ensure identical initial states
+    p_wires = p_kicker.copy()
+
+    # ----------------------------------------------------
+    # 5. Execute Tracking
+    # ----------------------------------------------------
+
+    # Case A: Single pass through NonlinearKicker
+    kicker.track(p_kicker)
+
+    # Case B: Sequential passes through the 3 Wires
+    for w in wire_elements:
+        w.track(p_wires)
+
+    # ----------------------------------------------------
+    # 6. Validation
+    # ----------------------------------------------------
+    # Use xo.assert_allclose to automatically handle GPU->CPU data transfers.
+    # We expect results to be nearly identical (double-precision floating point error margin).
+
+    xo.assert_allclose(p_kicker.px, p_wires.px, rtol=1e-13, atol=1e-13)
+
+    xo.assert_allclose(p_kicker.py, p_wires.py, rtol=1e-13, atol=1e-13)
+
+
+def test_input_validation():
+    # Pure CPU test to check if the __init__ method catches errors correctly
+
+    with pytest.raises(ValueError, match="Dimension mismatch"):
+        # Intentionally passing inconsistent array lengths
+        xt.NonlinearKicker(
+            current=[100, 200],  # Length 2
+            xma=[0.1],  # Length 1 (Incorrect!)
+            yma=[0.1, 0.2],
+        )

xtrack/beam_elements/elements.py

@@ -4222,3 +4222,47 @@ def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
 class ThinSliceNotNeededError(Exception):
     pass
 
+
+class NonlinearKicker(BeamElement):
+    """Beam element modeling nonlinear kicker composed of multiple wires (used for beam injection).
+
+    Parameters
+    ----------
+
+    L_phy : float
+        Physical length of the wires in meters. Assumed common for all wires. Default is ``0``.
+    L_int : float
+        Interaction length of the wires in meters. Assumed common for all wires. Default is ``0``.
+    current : array-like
+        Array of currents for each wire in Ampere. Default is ``0``.
+    xma : array-like
+        Array of horizontal positions for each wire in meters. Default is ``0``.
+    yma : array-like
+        Array of vertical positions for each wire in meters. Default is ``0``.
+    """
+
+    _xofields = {
+        "L_phy": xo.Float64,
+        "L_int": xo.Float64,
+        "current": xo.Float64[:],
+        "xma": xo.Float64[:],
+        "yma": xo.Float64[:],
+    }
+
+    _extra_c_sources = [
+        "#include <beam_elements/elements_src/nonlinearkicker.h>",
+    ]
+
+    def __init__(self, **kwargs):
+
+        # Validation: Ensure all arrays have the same length
+        c_len = len(kwargs.get("current", []))
+        x_len = len(kwargs.get("xma", []))
+        y_len = len(kwargs.get("yma", []))
+
+        if not (c_len == x_len == y_len):
+            raise ValueError(
+                f"Dimension mismatch: current ({c_len}), xma ({x_len}), and yma ({y_len}) must be equal."
+            )
+
+        super().__init__(**kwargs)

xtrack/beam_elements/elements_src/nonlinearkicker.h

@@ -0,0 +1,72 @@
+// copyright ############################### //
+// This file is part of the Xtrack Package.  //
+// Copyright (c) CERN, 2021.                 //
+// ######################################### //
+#ifndef XTRACK_NONLINEARKICKER_H
+#define XTRACK_NONLINEARKICKER_H
+
+#include <headers/track.h>
+
+
+GPUFUN
+void NonlinearKicker_track_local_particle(NonlinearKickerData el, LocalParticle* part0){
+
+
+    // Data from nonlinear kicker
+    double const L_phy = NonlinearKickerData_get_L_phy(el);
+    double const L_int = NonlinearKickerData_get_L_int(el); 
+
+    // Get the number of wires defined in the arrays
+    int64_t const n_wires = NonlinearKickerData_len_xma(el);
+
+    START_PER_PARTICLE_BLOCK(part0, part);
+
+        // Retrieve particle coordinates
+        double x      = LocalParticle_get_x(part);
+        double y      = LocalParticle_get_y(part);
+
+        // chi = q/q0 * m0/m
+        // p0c : reference particle momentum
+        // q0  : reference particle charge
+        //double const chi    = LocalParticle_get_chi(part);
+        double const p0c    = LocalParticle_get_p0c(part);
+        double const q0     = LocalParticle_get_q0(part);
+
+        double dpx_total = 0.0;
+        double dpy_total = 0.0;
+
+        // Loop over each wire
+        for (int64_t i = 0; i < n_wires; i++){
+
+            // Retrieve parameters for the i-th wire from the data arrays
+            double const cx = NonlinearKickerData_get_xma(el, i);
+            double const cy = NonlinearKickerData_get_yma(el, i);
+            double const curr = NonlinearKickerData_get_current(el, i);
+
+            // Calculate distance relative to the wire center
+            double D_x    = x-cx;
+            double D_y    = y-cy;
+            double R2     = D_x*D_x + D_y*D_y;
+
+            // Skip the calculation if the distance is too small to avoid division by zero
+            if (R2 < 1e-20) continue;
+
+            // Computing the kick
+            double const L1   = L_int + L_phy;
+            double const L2   = L_int - L_phy;
+            double const N    = MU_0*curr*q0/(4*PI*p0c/C_LIGHT);
+
+            double dpx_i  =  -N*D_x*(sqrt(L1*L1 + 4.0*R2) - sqrt(L2*L2 + 4.0*R2))/R2;
+            double dpy_i  =  -N*D_y*(sqrt(L1*L1 + 4.0*R2) - sqrt(L2*L2 + 4.0*R2))/R2;
+
+            dpx_total += dpx_i;
+            dpy_total += dpy_i;
+        }
+
+        // Update the particle properties
+        LocalParticle_add_to_px(part, dpx_total);
+        LocalParticle_add_to_py(part, dpy_total);
+    END_PER_PARTICLE_BLOCK;
+}
+
+#endif