Feature/longitudinal exciter with arbitrary digital waveforms

tests/test_elements.py

@@ -45,6 +45,8 @@ def test_constructor(test_context):
         xt.Wire(_context=test_context, current=3.),
         xt.Exciter(_context=test_context, knl=[1], samples=[1,2,3,4],
                    sampling_frequency=1e3),
+        xt.LongitudinalExciter(_context=test_context, voltage=1000., samples=[1,2,3,4],
+                   sampling_frequency=1e3),
         xt.Bend(_context=test_context, length=1.),
         xt.Quadrupole(_context=test_context, length=1.),
         xt.ElectronCooler(_context=test_context,current=2.4,length=1.5,radius_e_beam=25*1e-3,
@@ -140,6 +142,8 @@ def test_backtrack(test_context):
         xt.Elens(_context=test_context, inner_radius=0.1),
         xt.Exciter(_context=test_context, knl=[1], samples=[1,2,3],
                    sampling_frequency=1e3),
+        xt.LongitudinalExciter(_context=test_context, voltage=1000., samples=[1,2,3,4],
+                   sampling_frequency=1e3),
     ]
 
     dtk_particle = dtk.TestParticles(

tests/test_longitudinal_exciter.py

@@ -0,0 +1,77 @@
+import pathlib
+
+import numpy as np
+import xobjects as xo
+import xtrack as xt
+from xobjects.test_helpers import for_all_test_contexts
+
+test_data_folder = pathlib.Path(
+    __file__).parent.joinpath('../test_data').absolute()
+
+@for_all_test_contexts
+def test_longitudinal_exciter(test_context):
+    """
+    This test checks that tracking with a LongitudinalExciter element produces the same effect
+    as tracking with a cavity set to the same frequency, voltage, and phase, for off-momentum
+    particles. This ensures that the LongitudinalExciter is correctly implemented and consistent
+    with the established Cavity element for equivalent excitation.
+    """
+
+    line = xt.Line.from_json(test_data_folder /
+                            'hllhc15_thick/lhc_thick_with_knobs.json')
+    line.build_tracker(_context=test_context)
+
+    tw = line.twiss(method="4d")
+    line_exciter = line.copy()
+    line_cavity = line.copy()
+
+    num_turns = 1000
+    harmonic = 35640
+    voltage = 10e6
+    lag = 180
+    cav_dpp = 1e-3
+    t0 = tw['T_rev0']
+    f0 = harmonic/t0
+    cav_df = - cav_dpp*tw['slip_factor']*f0
+    cav_f = f0 + cav_df
+    sampling_freq = 100_000*cav_f
+
+    cavity = xt.Cavity(frequency=cav_f, voltage=voltage, lag=lag, absolute_time=True)
+
+    tarray = np.arange(0, 1/cav_f, 1/sampling_freq)
+    samples = np.sin(2*np.pi*cav_f*tarray + lag/180*np.pi)
+
+    longi_exciter = xt.LongitudinalExciter(voltage=voltage,
+                                        samples=samples,
+                                        sampling_frequency=sampling_freq,
+                                        frev=1/t0,
+                                        duration=num_turns*t0,
+                                        start_turn=0)
+
+    line_exciter.insert("exciter", longi_exciter, at=0.)
+    line_cavity.insert("cavity", cavity, at=0.)
+
+    part_exciter = line.build_particles(
+        method="4d",
+        zeta=0,
+        delta=np.linspace(-1e-3, 1e-3, 10),
+        x_norm=0,
+        px_norm=0,
+        y_norm=0,
+        py_norm=0,
+        nemitt_x=0,
+        nemitt_y=0,
+    )
+
+    part_cavity = part_exciter.copy()
+
+    line_exciter.track(part_exciter, num_turns=100, turn_by_turn_monitor=True)
+    line_cavity.track(part_cavity, num_turns=100, turn_by_turn_monitor=True)
+
+    delta_exciter = line_exciter.record_last_track.delta
+    zeta_exciter = line_exciter.record_last_track.zeta
+    delta_cavity = line_cavity.record_last_track.delta
+    zeta_cavity = line_cavity.record_last_track.zeta
+
+    xo.assert_allclose(delta_exciter, delta_cavity, atol=1e-5)
+    xo.assert_allclose(zeta_exciter, zeta_cavity, atol=1e-5)

xtrack/beam_elements/__init__.py

@@ -5,6 +5,7 @@
 
 from .elements import *
 from .exciter import Exciter
+from .longitudinal_exciter import LongitudinalExciter
 from .apertures import *
 from .magnets import Magnet
 from .beam_interaction import BeamInteraction, ParticlesInjectionSample

xtrack/beam_elements/elements_src/longitudinal_exciter.h

@@ -0,0 +1,62 @@
+// ##################################
+// LongitudinalExciter element
+//
+// Author: Pablo Arrutia
+// Date: 15.07.2025
+// ##################################
+
+#ifndef XTRACK_LONGITUDINAL_EXCITER_H
+#define XTRACK_LONGITUDINAL_EXCITER_H
+
+#include <headers/track.h>
+
+
+GPUFUN
+void LongitudinalExciter_track_local_particle(LongitudinalExciterData el, LocalParticle* part0){
+
+    // get parameters
+    double const voltage = LongitudinalExciterData_get_voltage(el);
+    GPUGLMEM float const* samples = LongitudinalExciterData_getp1_samples(el, 0);
+    int64_t const nsamples = LongitudinalExciterData_get_nsamples(el);
+	int64_t const nduration = LongitudinalExciterData_get_nduration(el);
+    double const sampling_frequency = LongitudinalExciterData_get_sampling_frequency(el);
+    double const frev = LongitudinalExciterData_get_frev(el);
+    int64_t const start_turn = LongitudinalExciterData_get_start_turn(el);
+
+    #ifdef XSUITE_BACKTRACK
+        #define XTRACK_LONGITUDINAL_EXCITER_SIGN (-1)
+    #else
+        #define XTRACK_LONGITUDINAL_EXCITER_SIGN (+1)
+    #endif
+
+    START_PER_PARTICLE_BLOCK(part0, part);
+        // zeta is the absolute path length deviation from the reference particle: zeta = (s - beta0*c*t)
+        // but without limits, i.e. it can exceed the circumference (for coasting beams)
+        // as the particle falls behind or overtakes the reference particle
+        double const zeta = LocalParticle_get_zeta(part);
+        double const at_turn = LocalParticle_get_at_turn(part);
+        double const beta0 = LocalParticle_get_beta0(part);
+
+        // compute excitation sample index
+        int64_t i = sampling_frequency * ( ( at_turn - start_turn ) / frev - zeta / beta0 / C_LIGHT );
+
+        if (i >= 0 && i < nduration){
+			if (i >= nsamples){
+				i = i % nsamples;
+			}
+
+            // scale voltage by excitation strength
+            double const scaled_voltage = voltage * samples[i];
+
+            // apply longitudinal kick (energy change)
+            double const q = fabs(LocalParticle_get_q0(part)) * LocalParticle_get_charge_ratio(part);
+            double const energy = XTRACK_LONGITUDINAL_EXCITER_SIGN * q * scaled_voltage;
+
+            // Apply energy change to particle
+            LocalParticle_add_to_energy(part, energy, 1);
+
+        }
+    END_PER_PARTICLE_BLOCK;
+}
+
+#endif

xtrack/beam_elements/longitudinal_exciter.py

@@ -0,0 +1,120 @@
+"""
+LongitudinalExciter element
+
+Author: Pablo Arrutia
+Date: 15.07.2025
+
+"""
+
+import xobjects as xo
+import numpy as np
+
+from ..base_element import BeamElement
+from ..general import _pkg_root
+
+
+class LongitudinalExciter(BeamElement):
+    """Beam element modeling a longitudinal exciter as a time-dependent voltage source. 
+
+    The given voltage is scaled according to a custom waveform,
+    allowing for arbitrary time dependence. The waveform is specified by an array of samples:
+
+        voltage(t) = voltage * samples[i]
+
+    It is *not* assumed that the variations are slow compared to the revolution frequency
+    and the particle arrival time is taken into account when determining the sample index:
+
+        i = sampling_frequency * ( ( at_turn - start_turn ) / f_rev - zeta / beta0 / c0 )
+
+    where zeta=(s-beta0*c0*t) is the longitudinal coordinate of the particle, beta0 the
+    relativistic beta factor of the particle, c0 is the speed of light, at_turn is the
+    current turn number, f_rev is the revolution frequency, and sampling_frequency is the sampling
+    frequency. The excitation starts with the first sample when the reference particle 
+    arrives at the element in start_turn.
+
+    For example, to compute samples for a sinusoidal excitation with frequency f_ex one
+    would calculate the waveform as: samples[i] = np.sin(2*np.pi*f_ex*i/sampling_frequency)
+
+    Notes:
+        - This is similar to an Exciter but applies longitudinal kicks instead of transverse kicks.
+        - The voltage is applied as an energy change to the particles, similar to a Cavity.
+
+    Parameters:
+        - voltage (float): Base voltage in Volts that will be scaled by the waveform.
+        - samples (float array): Samples of excitation strength to scale voltage as function of time.
+        - nsamples (int): Number of samples. Pass this instead of samples to reserve memory for later initialisation.
+        - sampling_frequency (float): Sampling frequency in Hz.
+        - frev (float): Revolution frequency in Hz of circulating beam (used to relate turn number to sample index).
+        - start_turn (int): Turn of the reference particle when to start excitation.
+        - duration (float): Duration of excitation in s (defaults to nsamples/sampling_frequency). Repeats the waveform to fill the duration.
+
+    Example:
+        >>> fs = 10e6 # sampling frequency in Hz
+        >>> 
+        >>> # A simple sine at 500 kHz ...
+        >>> t = np.arange(1000)/fs
+        >>> f_ex = 5e5 # excitation frequency in Hz
+        >>> signal = np.sin(2*np.pi*f_ex*t)
+        >>> 
+        >>> # create the longitudinal exciter
+        >>> frev = 1e6 # revolution frequency in Hz
+        >>> voltage = 1000 # this is scaled by the waveform
+        >>> long_exciter = LongitudinalExciter(samples=signal, sampling_frequency=fs, frev=frev, start_turn=0, voltage=voltage)
+        >>> 
+        >>> # add it to the line
+        >>> line.insert_element(index=..., name=..., element=long_exciter)
+
+    """
+
+    _xofields={
+        'voltage': xo.Float64,
+        'nsamples': xo.Int64,
+        'sampling_frequency': xo.Float64,
+        'frev': xo.Float64,
+        'start_turn': xo.Int64,
+        'nduration': xo.Int64,
+        'samples': xo.Float32[:],
+        }
+
+    has_backtrack = True
+
+    _extra_c_sources = ['#include <beam_elements/elements_src/longitudinal_exciter.h>']
+
+
+    def __init__(self, *, samples=None, nsamples=None, sampling_frequency=0., frev=0., voltage=0., start_turn=0, duration=None, _xobject=None, **kwargs):
+
+        if _xobject is not None:
+            super().__init__(_xobject=_xobject)
+
+        else:
+            # determine sample length or create empty samples for later initialisation
+            if samples is not None:
+                if nsamples is not None and nsamples != len(samples):
+                    raise ValueError("Only one of samples or nsamples may be specified")
+                nsamples = len(samples)
+            if samples is None:
+                samples = np.zeros(nsamples)
+            kwargs["nduration"] = nsamples if duration is None else (duration * sampling_frequency)
+
+            super().__init__(
+                voltage=voltage, samples=samples, nsamples=nsamples,
+                sampling_frequency=sampling_frequency,
+                frev=frev, start_turn=start_turn, **kwargs)
+
+    @property
+    def duration(self):
+        return self.nduration / self.sampling_frequency
+
+    @duration.setter
+    def duration(self, duration):
+        self.nduration = duration * self.sampling_frequency
+
+    def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
+        ctx2np = self._buffer.context.nparray_from_context_array
+        return self.__class__(voltage=-ctx2np(self.voltage),
+                              samples=self.samples,
+                              nsamples=self.nsamples,
+                              sampling_frequency=self.sampling_frequency,
+                              frev=self.frev,
+                              start_turn=self.start_turn,
+                              _context=_context, _buffer=_buffer, _offset=_offset)