Add time-dependent noise handling inside C implementation of Elens beam element

tests/test_elements.py

@@ -361,6 +361,122 @@ def compute_coef(r_measured, j_measured, r_1_new, r_2_new,
                       rtol=1e-2, atol=1e-2)
 
 
+@for_all_test_contexts
+def test_elens_noise_pattern(test_context):
+    # Set fixed numpy seed
+    np.random.seed(42)
+    # Create a binary noise pattern of 25 samples
+    noise_pattern = np.random.randint(0, 2, 25)
+
+    line_no_pattern = xt.Line(
+        elements=[
+            xt.Drift(length=1e-3),
+            xt.Elens(
+                inner_radius=1.1e-3,
+                outer_radius=2.2e-3,
+                elens_length=3.0,
+                voltage=15e3,
+                current=5,
+            ),
+        ]
+    )
+
+    noise_modes = ["loop", "zeros", "ones", "last"]
+    lines_with_patterns = {
+        mode: xt.Line(
+            elements=[
+                xt.Drift(length=1e-3),
+                xt.Elens(
+                    inner_radius=1.1e-3,
+                    outer_radius=2.2e-3,
+                    elens_length=3.0,
+                    voltage=15e3,
+                    current=5,
+                    noise=noise_pattern,
+                    noise_mode=mode,
+                ),
+            ]
+        )
+        for mode in noise_modes
+    }
+
+    # Build trackers for all lines
+    line_no_pattern.build_tracker(_context=test_context)
+    for mode, line in lines_with_patterns.items():
+        line.build_tracker(_context=test_context)
+
+    # Test each noise mode
+    for mode, line in lines_with_patterns.items():
+        particle_ref = xp.Particles(
+            p0c=np.array([7000e9]),
+            x=np.array([1e-3]),
+            px=np.array([0.0]),
+            y=np.array([2.2e-3]),
+            py=np.array([0.0]),
+            zeta=np.array([0.0]),
+            context=test_context,
+        )
+        particle_ref_copy = particle_ref.copy(_context=test_context)
+
+        if mode == "loop":
+            # Explicitly track with the line_no_pattern for comparison
+            for i in range(50):
+                line_no_pattern.element_dict["e1"].current = 5 * noise_pattern[i % 25]
+                line_no_pattern.track(particle_ref)
+
+        elif mode == "zeros":
+            # All values after the array length are zero
+            for i in range(50):
+                current_value = 5 * (noise_pattern[i] if i < len(noise_pattern) else 0)
+                line_no_pattern.element_dict["e1"].current = current_value
+                line_no_pattern.track(particle_ref)
+
+        elif mode == "ones":
+            # All values after the array length are one
+            for i in range(50):
+                current_value = 5 * (noise_pattern[i] if i < len(noise_pattern) else 1)
+                line_no_pattern.element_dict["e1"].current = current_value
+                line_no_pattern.track(particle_ref)
+
+        elif mode == "last":
+            # All values after the array length are the last array value
+            for i in range(50):
+                current_value = 5 * (
+                    noise_pattern[i] if i < len(noise_pattern) else noise_pattern[-1]
+                )
+                line_no_pattern.element_dict["e1"].current = current_value
+                line_no_pattern.track(particle_ref)
+
+        # Track the line with the specified noise mode
+        line.track(particle_ref_copy, num_turns=50)
+
+        # Validate results (all particles must be the same)
+        xo.assert_allclose(
+            test_context.nparray_from_context_array(particle_ref.x)[0],
+            particle_ref_copy.x,
+            rtol=1e-14,
+            atol=1e-14,
+        )
+        xo.assert_allclose(
+            test_context.nparray_from_context_array(particle_ref.y)[0],
+            particle_ref_copy.y,
+            rtol=1e-14,
+            atol=1e-14,
+        )
+        xo.assert_allclose(
+            test_context.nparray_from_context_array(particle_ref.px)[0],
+            particle_ref_copy.px,
+            rtol=1e-14,
+            atol=1e-14,
+        )
+        xo.assert_allclose(
+            test_context.nparray_from_context_array(particle_ref.py)[0],
+            particle_ref_copy.py,
+            rtol=1e-14,
+            atol=1e-14,
+        )
+
+
 @for_all_test_contexts
 def test_wire(test_context):
     dtk_particle = dtk.TestParticles(

xtrack/beam_elements/elements.py

@@ -172,7 +172,19 @@ class Elens(BeamElement):
         Array of coefficients of the polynomial. Default is ``[0]``.
     polynomial_order : int
         Order of the polynomial. Default is ``0``.
-
+    noise: array
+        Array of noise values. Default is ``[1.0]``. Use this to model different
+        on/off states of the elens for each turn. The length of the noise array
+        should be equal to the number of turns you want to simulate. Otherwise,
+        the noise_mode will be used to handle the noise array. 1.0 means full elens
+        and 0.0 means no elens.
+    noise_mode: {'loop', 'last', 'zeros', 'ones'}
+        How to handle the noise array if the number of turns is larger than the
+        length of the noise array. Default is ``'loop'``.
+        'loop': loop over the noise array.
+        'last': use the last value of the noise array.
+        'zeros': use zero (no elens) for the remaining turns.
+        'ones': use one (full elens) for the remaining turns.
     '''
 
     _xofields={
@@ -185,6 +197,9 @@ class Elens(BeamElement):
         'residual_kick_y': xo.Float64,
         'coefficients_polynomial': xo.Field(xo.Float64[:], default=[0]),
         'polynomial_order': xo.Float64,
+        'len_noise': xo.Int64,
+        'noise': xo.Field(xo.Float64[:], default=[1.]),
+        'noise_mode': xo.Int64,
     }
 
     has_backtrack = True
@@ -193,10 +208,20 @@ class Elens(BeamElement):
         _pkg_root.joinpath('beam_elements/elements_src/elens.h')]
 
     def __init__(self, _xobject=None, **kwargs):
+        kwargs["noise_mode"] = {"loop": 0, "last": 1, "zeros": 2, "ones": 3}.get(
+            kwargs.get("noise_mode", "loop"),
+            ValueError("noise_mode must be 'loop', 'last', 'zeros' or 'ones'")
+        )
+
+        if isinstance(kwargs["noise_mode"], ValueError):
+            raise kwargs["noise_mode"]
+
         super().__init__(_xobject=_xobject, **kwargs)
         if _xobject is None:
             polynomial_order = len(self.coefficients_polynomial) - 1
             self.polynomial_order = polynomial_order
+            len_noise = len(self.noise)
+            self.len_noise = len_noise
 
 
 class NonLinearLens(BeamElement):

xtrack/beam_elements/elements_src/elens.h

@@ -23,11 +23,34 @@ void Elens_track_local_particle(ElensData el, LocalParticle* part0){
     double const residual_kick_y = ElensData_get_residual_kick_y(el);
 
     int const polynomial_order = ElensData_get_polynomial_order(el);
+    int const len_noise = ElensData_get_len_noise(el);
 
     /*gpuglmem*/ double const* coefficients_polynomial =
                                 ElensData_getp1_coefficients_polynomial(el, 0);
 
     //start_per_particle_block (part0->part)
+        int at_turn = LocalParticle_get_at_turn(part);
+        const int noise_mode = ElensData_get_noise_mode(el);
+        double noise;
+        if (noise_mode == 0)
+        {
+            at_turn = (at_turn % len_noise + len_noise) % len_noise; // Handle wrapping for negative values
+            noise = ElensData_get_noise(el, at_turn);
+        }
+        else if (noise_mode == 1 || noise_mode == 2 || noise_mode == 3)
+        {
+            if (at_turn >= len_noise)
+                noise = (noise_mode == 1)   ? ElensData_get_noise(el, len_noise - 1)
+                        : (noise_mode == 2) ? 0.0
+                                            : 1.0; // noise_mode == 3
+            else
+                noise = ElensData_get_noise(el, at_turn);
+        }
+        else
+        {
+            // Error case: default to noise = 1.0
+            noise = 1.0;
+        }
 
         // electron mass
         double const EMASS  = 510998.928;
@@ -120,7 +143,9 @@ void Elens_track_local_particle(ElensData el, LocalParticle* part0){
 
         // # calculate the kick at r2 (maximum kick)
         double theta_max = ((1.0/(4.0*PI*EPSILON_0)));
-        theta_max = theta_max*(2*elens_length*current);
+
+        double actual_current = current*noise;
+        theta_max = theta_max * (2 * elens_length * actual_current);
 
         // for the moment: e-beam in the opposite direction from proton beam
         // generalize later