fix:'replace_spacecharge_with_PIC' now creates 'PICCollection' with longitudinal kick for 'FFTSolver3D'

Added 'apply_z_kick' parameter to 'replace_spacecharge_with_PIC'. Added logic so that if solver is 'FFTSolver3D' and no 'apply_z_kick' parameter value is specified, 'apply_z_kick' is set to True when creating 'PICCollection'. Previously, PICCollection would always be created with 'apply_z_kick' as False even when using 'FFTSolver3D', and therefore no longitudinal kick would be deposited.

Author: MarkusSQ

tests/test_replace_spacecharge_with_PIC_z_kick.py

@@ -0,0 +1,93 @@
+import pytest
+
+import xtrack as xt
+import xpart as xp
+import xfields as xf
+
+from xobjects.test_helpers import for_all_test_contexts
+from xfields.config_tools import replace_spacecharge_with_PIC
+
+
+@pytest.mark.parametrize(
+    "solver, user_apply, expected",
+    [
+        ("FFTSolver2p5D", None, False),
+        ("FFTSolver3D", None, True),
+        ("FFTSolver2p5D", True,  True),
+        ("FFTSolver3D", False, False),
+    ]
+)
+@for_all_test_contexts
+def test_replace_spacecharge_with_PIC_apply_z_kick_logic(
+        solver, user_apply, expected, test_context):
+    bunch_intensity = 2.5e11
+    sigma_x = 3e-3
+    sigma_y = 2e-3
+    sigma_z = 0.30
+    p0c = 25.92e9
+
+    lprof = xf.LongitudinalProfileQGaussian(
+        _context=test_context,
+        number_of_particles=bunch_intensity,
+        sigma_z=sigma_z,
+        z0=0.0,
+        q_parameter=1.0
+    )
+
+    sc1 = xf.SpaceChargeBiGaussian(
+        _context=test_context,
+        update_on_track=False,
+        length=1.0,
+        longitudinal_profile=lprof,
+        mean_x=0.0, mean_y=0.0,
+        sigma_x=sigma_x, sigma_y=sigma_y,
+        min_sigma_diff=1e-12,
+    )
+    sc2 = xf.SpaceChargeBiGaussian(
+        _context=test_context,
+        update_on_track=False,
+        length=2.0,
+        longitudinal_profile=lprof,
+        mean_x=0.0, mean_y=0.0,
+        sigma_x=sigma_x, sigma_y=sigma_y,
+        min_sigma_diff=1e-12,
+    )
+
+    line = xt.Line(elements=[
+        xt.Drift(length=0.1),
+        sc1,
+        xt.Drift(length=0.2),
+        sc2,
+        xt.Drift(length=0.1),
+    ])
+
+    line.build_tracker(_context=test_context, compile=False)
+    line.particle_ref = xp.Particles(_context=test_context, p0c=p0c)
+
+    n_sigmas_range_pic_x = 4.0
+    n_sigmas_range_pic_y = 4.0
+    nx_grid = 16
+    ny_grid = 16
+    nz_grid = 16
+    n_lims_x = 3
+    n_lims_y = 3
+    z_range = (-3 * sigma_z, 3 * sigma_z)
+
+    kwargs = dict(
+        line=line,
+        n_sigmas_range_pic_x=n_sigmas_range_pic_x,
+        n_sigmas_range_pic_y=n_sigmas_range_pic_y,
+        nx_grid=nx_grid, ny_grid=ny_grid, nz_grid=nz_grid,
+        n_lims_x=n_lims_x, n_lims_y=n_lims_y,
+        z_range=z_range,
+        solver=solver,
+    )
+    if user_apply is not None:
+        kwargs["apply_z_kick"] = user_apply
+
+    pic_collection, all_pics = replace_spacecharge_with_PIC(**kwargs)
+
+    assert pic_collection.apply_z_kick == expected, (
+        f"PICCollection.apply_z_kick mismatch for solver={solver}, "
+        f"user_apply={user_apply}; expected {expected}"
+    )

xfields/config_tools/spacecharge_config_tools.py

@@ -277,6 +277,7 @@ def replace_spacecharge_with_PIC(
         n_sigmas_range_pic_x, n_sigmas_range_pic_y,
         nx_grid, ny_grid, nz_grid, n_lims_x, n_lims_y, z_range,
         solver='FFTSolver2p5D',
+        apply_z_kick=None,
         _context=None,
         _buffer=None):
 
@@ -303,6 +304,10 @@ def replace_spacecharge_with_PIC(
         Number different limits in y for which PIC need to be generated.
     z_range : float
         Range of the longitudinal grid.
+    apply_z_kick : bool, optional
+        If ``True``, the longitudinal kick on the particles is applied.
+        Defaults to ``True`` when ``solver`` is ``FFTSolver3D``,
+        otherwise ``False``.
     _context : xtrack.Context (optional)
         Context in which the PIC elements are created.
     _buffer : xtrack.Buffer (optional)
@@ -338,6 +343,9 @@ def replace_spacecharge_with_PIC(
     y_lim_min = np.min(all_sigma_y) * (n_sigmas_range_pic_y - 0.5)
     y_lim_max = np.max(all_sigma_y) * (n_sigmas_range_pic_y + 0.5)
 
+    if apply_z_kick is None:
+        apply_z_kick = (solver == "FFTSolver3D")
+
     pic_collection = PICCollection(
         _context=_context,
         _buffer=_buffer,
@@ -346,6 +354,7 @@ def replace_spacecharge_with_PIC(
         y_lim_min=y_lim_min, y_lim_max=y_lim_max, n_lims_y=n_lims_y,
         z_range=z_range,
         solver=solver,
+        apply_z_kick=apply_z_kick,
         gamma0=line.particle_ref.gamma0[0])
 
     all_pics = []