Solenoid: apply radiation after multipolar kick

tests/test_elements_thick.py

@@ -1156,6 +1156,130 @@ def test_solenoid_with_mult_kicks(test_context, backtrack):
     xo.assert_allclose(p_test.pzeta, p_ref.pzeta, rtol=0, atol=1e-13)
 
 
+@for_all_test_contexts
+@pytest.mark.parametrize(
+    'radiation_mode,config',
+    [
+        ('mean', {}),
+        ('mean', {'XTRACK_SYNRAD_KICK_SAME_AS_FIRST': True}),
+        ('mean', {'XTRACK_SYNRAD_SCALE_SAME_AS_FIRST': True}),
+        ('quantum', {}),
+    ],
+)
+def test_drift_like_solenoid_with_kicks_radiation(test_context, radiation_mode, config):
+    config['XTRACK_USE_EXACT_DRIFTS'] = True
+    knl = [0.1, 0.4, 0.5]
+    ksl = [0.2, 0.3, 0.6]
+
+    line_test = xt.Line(elements=[
+        xt.Drift(length=0.5),
+        xt.Multipole(knl=knl, ksl=ksl),
+        xt.Drift(length=0.5),
+    ])
+
+    line_ref = xt.Line(elements=[
+        xt.Solenoid(ks=0, length=1, knl=knl, ksl=ksl, num_multipole_kicks=1)
+    ])
+
+    coords = np.linspace(-0.05, 0.05, 10)
+    coords_6d = np.array(list(itertools.product(*(coords,) * 6))).T
+
+    p0 = xt.Particles(
+        _context=test_context,
+        x=coords_6d[0],
+        px=coords_6d[1],
+        y=coords_6d[2],
+        py=coords_6d[3],
+        zeta=coords_6d[4],
+        delta=coords_6d[5],
+    )
+    p0._init_random_number_generator(seeds=np.arange(len(coords_6d[0]), dtype=int))
+    p_ref = p0.copy()
+    p_test = p0.copy()
+
+    line_ref.build_tracker(_context=test_context)
+    line_ref.configure_radiation(model=radiation_mode)
+    line_ref.config.update(config)
+    line_ref.track(p_ref, num_turns=1)
+
+    line_test.build_tracker(_context=test_context)
+    line_test.configure_radiation(model=radiation_mode)
+    line_test.config.update(config)
+    line_test.track(p_test, num_turns=1)
+
+    xo.assert_allclose(p_test.x, p_ref.x, rtol=0, atol=1e-13)
+    xo.assert_allclose(p_test.px, p_ref.px, rtol=0, atol=1e-13)
+    xo.assert_allclose(p_test.y, p_ref.y, rtol=0, atol=1e-13)
+    xo.assert_allclose(p_test.py, p_ref.py, rtol=0, atol=1e-13)
+    xo.assert_allclose(p_test.delta, p_ref.delta, rtol=0, atol=1e-13)
+    xo.assert_allclose(p_test.pzeta, p_ref.pzeta, rtol=0, atol=1e-13)
+
+
+@for_all_test_contexts
+@pytest.mark.parametrize(
+    'radiation_mode,config',
+    [
+        ('mean', {}),
+        ('mean', {'XTRACK_SYNRAD_KICK_SAME_AS_FIRST': True}),
+        ('mean', {'XTRACK_SYNRAD_SCALE_SAME_AS_FIRST': True}),
+        ('quantum', {}),
+    ],
+)
+def test_solenoid_with_kicks_radiation(test_context, radiation_mode, config):
+    config['XTRACK_USE_EXACT_DRIFTS'] = True
+
+    ks = 0.4
+    l = 1.1
+    knl = [0.1, 0.4, 0.5]
+    ksl = [0.2, 0.3, 0.6]
+
+    sol_ref = xt.Solenoid(ks=ks, length=l)
+    sol_1 = xt.Solenoid(ks=ks, length=l, knl=knl, ksl=ksl, num_multipole_kicks=1)
+    sol_3 = xt.Solenoid(ks=ks, length=l, knl=knl, ksl=ksl, num_multipole_kicks=3)
+
+    line_ref = xt.Line(elements=[sol_ref])
+    line_1 = xt.Line(elements=[sol_1])
+    line_3 = xt.Line(elements=[sol_3])
+
+    coords = np.linspace(-0.05, 0.05, 10)
+    coords_6d = np.array(list(itertools.product(*(coords,) * 6))).T
+
+    p0 = xt.Particles(
+        _context=test_context,
+        x=coords_6d[0],
+        px=coords_6d[1],
+        y=coords_6d[2],
+        py=coords_6d[3],
+        zeta=coords_6d[4],
+        delta=coords_6d[5],
+    )
+    p0._init_random_number_generator(seeds=np.arange(len(coords_6d[0]), dtype=int))
+    p_ref = p0.copy()
+    p_1 = p0.copy()
+    p_3 = p0.copy()
+
+    line_ref.build_tracker(_context=test_context)
+    line_ref.track(p_ref, num_turns=1)
+
+    line_1.build_tracker(_context=test_context)
+    line_1.configure_radiation(model=radiation_mode)
+    line_1.config.update(config)
+    line_1.track(p_1, num_turns=1)
+
+    line_3.build_tracker(_context=test_context)
+    line_3.configure_radiation(model=radiation_mode)
+    line_3.config.update(config)
+    line_3.track(p_3, num_turns=1)
+
+    d_delta_ref = p_ref.delta - p0.delta
+    d_delta_1 = p_1.delta - p0.delta
+    d_delta_3 = p_3.delta - p0.delta
+
+    xo.assert_allclose(
+        d_delta_1 - d_delta_ref, d_delta_3 - d_delta_ref, rtol=0.01, atol=1e-15
+    )
+
+
 @for_all_test_contexts
 def test_skew_quadrupole(test_context):
     k1 = 1.0

tests/test_slice_elements.py

@@ -1,8 +1,11 @@
-import xtrack as xt
+import itertools
+
+import numpy as np
 import pytest
 
-from xobjects.test_helpers import for_all_test_contexts
 import xobjects as xo
+import xtrack as xt
+from xobjects.test_helpers import for_all_test_contexts
 
 assert_allclose= xo.assert_allclose
 
@@ -630,9 +633,9 @@ def test_thick_slice_octupole(test_context):
 @for_all_test_contexts
 def test_thick_slice_solenoid(test_context):
 
-    oct = xt.Solenoid(ks=0.1, length=1)
+    sol = xt.Solenoid(ks=0.1, length=1)
 
-    line = xt.Line(elements=[oct])
+    line = xt.Line(elements=[sol])
 
     line.slice_thick_elements(
         slicing_strategies=[xt.Strategy(xt.Teapot(10, mode='thick'))])
@@ -646,7 +649,7 @@ def test_thick_slice_solenoid(test_context):
     p_slice = p0.copy()
 
     line.track(p_slice)
-    oct.track(p_ref)
+    sol.track(p_ref)
 
     assert_allclose(p_slice.x, p_ref.x, rtol=0, atol=1e-10)
     assert_allclose(p_slice.px, p_ref.px, rtol=0, atol=1e-10)

tests/test_slicing.py

@@ -280,7 +280,7 @@ def test_slicing_strategy_matching():
         if name == 'keep_this':
             assert isinstance(element, xt.Quadrupole)
         elif name == 'keep_drifts':
-            assert (element, xt.Drift)
+            assert isinstance(element, xt.Drift)
         elif name.startswith('drift_'):
             assert 'DriftSlice' in type(element).__name__
         elif name == 'keep_thin':

tests/test_spacecharge_in_ring.py

@@ -157,7 +157,7 @@ def test_ring_with_spacecharge(test_context, mode):
     else:
         assert isinstance(line[0], xf.SpaceChargeBiGaussian)
 
-    if mode is not 'frozen':
+    if mode != 'frozen':
         assert line.iscollective
     else:
         assert not line.iscollective

xtrack/beam_elements/elements_src/solenoid.h

@@ -1,6 +1,6 @@
 // copyright ############################### //
 // This file is part of the Xtrack Package.  //
-// Copyright (c) CERN, 2023.                 //
+// Copyright (c) CERN, 2024.                 //
 // ######################################### //
 
 #ifndef XTRACK_SOLENOID_H
@@ -20,13 +20,13 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) {
         factor_knl_ksl = -1;
     #endif
 
-    #ifndef XTRACK_SOLENOID_NO_SYNRAD  // does this do the right thing?
-    double dp_record_entry = 0.;
-    double dpx_record_entry = 0.;
-    double dpy_record_entry = 0.;
-    double dp_record_exit = 0.;
-    double dpx_record_exit = 0.;
-    double dpy_record_exit = 0.;
+    #ifndef XTRACK_SOLENOID_NO_SYNRAD
+        double dp_record_entry = 0.;
+        double dpx_record_entry = 0.;
+        double dpy_record_entry = 0.;
+        double dp_record_exit = 0.;
+        double dpx_record_exit = 0.;
+        double dpy_record_exit = 0.;
     #endif
 
     int64_t num_multipole_kicks = SolenoidData_get_num_multipole_kicks(el);
@@ -37,22 +37,34 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) {
     const double slice_length = length / (num_multipole_kicks + 1);
     const double kick_weight = 1. / num_multipole_kicks;
 
+    //start_per_particle_block (part0->part)
+    #ifndef XTRACK_SOLENOID_NO_SYNRAD
+        double const old_px = LocalParticle_get_px(part);
+        double const old_py = LocalParticle_get_py(part);
+        double const old_ax = LocalParticle_get_ax(part);
+        double const old_ay = LocalParticle_get_ay(part);
+        double const old_zeta = LocalParticle_get_zeta(part);
+    #endif
+
     for (int ii = 0; ii < num_multipole_kicks; ii++) {
-        //start_per_particle_block (part0->part)
-        Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag,
-                        &dp_record_entry, &dpx_record_entry, &dpy_record_entry,
-                        &dp_record_exit, &dpx_record_exit, &dpy_record_exit);
+        Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag);
 
         track_multipolar_kick_bend(
                     part, order, inv_factorial_order, knl, ksl, factor_knl_ksl,
                     kick_weight, 0, 0, 0, 0);
-        //end_per_particle_block
     }
 
-    //start_per_particle_block (part0->part)
-    Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag,
-                    &dp_record_entry, &dpx_record_entry, &dpy_record_entry,
-                    &dp_record_exit, &dpx_record_exit, &dpy_record_exit);
+    Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag);
+
+    #ifndef XTRACK_SOLENOID_NO_SYNRAD
+        if (radiation_flag > 0 && length > 0){
+            Solenoid_apply_radiation_single_particle(
+                part, length, radiation_flag,
+                old_px, old_py, old_ax, old_ay, old_zeta,
+                &dp_record_entry, &dpx_record_entry, &dpy_record_entry
+            );
+        }
+    #endif
     //end_per_particle_block
 }
 

xtrack/beam_elements/elements_src/thick_slice_bend.h

@@ -19,8 +19,8 @@ void ThickSliceBend_track_local_particle(
     const double h = ThickSliceBendData_get__parent_h(el);
     const double order = ThickSliceBendData_get__parent_order(el);
     const double inv_factorial_order = ThickSliceBendData_get__parent_inv_factorial_order(el);
-    const double* knl = ThickSliceBendData_getp1__parent_knl(el, 0);
-    const double* ksl = ThickSliceBendData_getp1__parent_ksl(el, 0);
+    /*gpuglmem*/ const double* knl = ThickSliceBendData_getp1__parent_knl(el, 0);
+    /*gpuglmem*/ const double* ksl = ThickSliceBendData_getp1__parent_ksl(el, 0);
     const int64_t model = ThickSliceBendData_get__parent_model(el);
 
     const int64_t num_multipole_kicks_parent = ThickSliceBendData_get__parent_num_multipole_kicks(el);

xtrack/beam_elements/elements_src/thick_slice_quadrupole.h

@@ -19,8 +19,8 @@ void ThickSliceQuadrupole_track_local_particle(
     const int64_t num_multipole_kicks_parent = ThickSliceQuadrupoleData_get__parent_num_multipole_kicks(el);
     const double order = ThickSliceQuadrupoleData_get__parent_order(el);
     const double inv_factorial_order = ThickSliceQuadrupoleData_get__parent_inv_factorial_order(el);
-    const double* knl = ThickSliceQuadrupoleData_getp1__parent_knl(el, 0);
-    const double* ksl = ThickSliceQuadrupoleData_getp1__parent_ksl(el, 0);
+    /*gpuglmem*/ const double* knl = ThickSliceQuadrupoleData_getp1__parent_knl(el, 0);
+    /*gpuglmem*/ const double* ksl = ThickSliceQuadrupoleData_getp1__parent_ksl(el, 0);
 
     #ifndef XSUITE_BACKTRACK
         double const length = weight * ThickSliceQuadrupoleData_get__parent_length(el); // m

xtrack/beam_elements/elements_src/thick_slice_solenoid.h

@@ -24,10 +24,9 @@ void ThickSliceSolenoid_track_local_particle(
 
 
     //start_per_particle_block (part0->part)
-        Solenoid_thick_track_single_particle(
+        Solenoid_thick_with_radiation_track_single_particle(
                 part, length, ks,
                 0, // radiation flag, not supported for now
-                NULL, NULL, NULL,
                 NULL, NULL, NULL);
     //end_per_particle_block
 

xtrack/beam_elements/elements_src/track_bend.h

@@ -19,7 +19,7 @@
 void Bend_track_local_particle_from_params(LocalParticle* part0,
                 double length, double k0, double k1, double h,
                 int64_t num_multipole_kicks, int64_t model,
-                double const* knl, double const* ksl,
+                /*gpuglmem*/ double const* knl, /*gpuglmem*/ double const* ksl,
                 int64_t order, double inv_factorial_order,
                 double factor_knl_ksl){
 

xtrack/beam_elements/elements_src/track_multipolar_components.h

@@ -9,8 +9,8 @@
 /*gpufun*/
 void track_multipolar_kick_bend(
     LocalParticle* part, int64_t order, double inv_factorial_order,
-    const double* knl,
-    const double* ksl,
+    /*gpuglmem*/ const double* knl,
+    /*gpuglmem*/ const double* ksl,
     double const factor_knl_ksl,
     double kick_weight, double k0, double k1, double h, double length){
 

xtrack/beam_elements/elements_src/track_quadrupole.h

@@ -35,7 +35,7 @@ void normal_quad_with_rotation_track(
 void Quadrupole_from_params_track_local_particle(
         double length, double k1, double k1s,
         int64_t num_multipole_kicks,
-        double const* knl, double const* ksl,
+        /*gpuglmem*/ double const* knl, /*gpuglmem*/ double const* ksl,
         int64_t order, double inv_factorial_order,
         double factor_knl_ksl,
         LocalParticle* part0