diff --git a/tests/test_nonlinearkicker.py b/tests/test_nonlinearkicker.py new file mode 100644 index 000000000..c9d440aec --- /dev/null +++ b/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], + ) diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 4ac4d505e..a796ff7c1 100644 --- a/xtrack/beam_elements/elements.py +++ b/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 ", + ] + + 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) \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/nonlinearkicker.h b/xtrack/beam_elements/elements_src/nonlinearkicker.h new file mode 100644 index 000000000..d53f234ec --- /dev/null +++ b/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 + + +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