Momentum Kick in the 2.5D Gaussian Space Charge Model

The test compares the space charge momentum kick experienced by particles within a long, initially cold, Gaussian bunch against the 2.5D analytical model in:

  1. Qiang, Phys. Rev. Accel. Beams 28, 114602 (2025), eqs. (31-32).

We use a cold (zero emittance) 100 MeV electron bunch in a drift space, with a single space charge slice.

Here \(\sigma_x = \sigma_y\), so the predicted momentum kick can be evaluated in closed form.

In this test, the final momentum of every particle is compared against its predicted value. Both the maximum error and its rms value over the bunch population are returned. Both values must be zero within a small tolerance.

The test can be run using either the Gauss2p5D space charge model or the 2.5D PIC space charge model.

Run

This example can be run as:

  • Python script: python3 run_sc_kick_Gauss2p5D.py or python3 run_sc_kick_PIC2p5D.py

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 259 You can copy this file from examples/space_charge_gaussian_kick/run_sc_kick_Gauss2p5D.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell, Ji Qiang
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = "Gauss2p5D"
sim.space_charge_gauss_nint = 101
sim.space_charge_gauss_charge_z_bins = 129
sim.space_charge_gauss_taylor_delta = 0.01
sim.space_charge_gauss_long_scale = 6.0
sim.slice_step_diagnostics = True

# domain decomposition & space charge mesh
sim.init_grids()

# beam energy ad bunch charge
kin_energy_MeV = 100  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 100000  # number of macro particles

# intialize reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

# initialize particle bunch
distr = distribution.Gaussian(
    lambdaX=4.0e-5,
    lambdaY=4.0e-5,
    lambdaT=1.0e-3,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice)
ns = 1  # number of slices per ds in the element
L = 1.0  # drift length
sim.lattice.extend([monitor, elements.Drift(name="d1", ds=L, nslice=ns), monitor])

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 260 You can copy this file from examples/space_charge_gaussian_kick/run_sc_kick_PIC2p5D.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell, Ji Qiang
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = "2p5D"
sim.poisson_solver = "fft"
sim.space_charge_num_longitudinal_bins = 129
sim.space_charge_apply_longitudinal_kick = True
sim.slice_step_diagnostics = True

# set numerical parameters and IO control
sim.max_level = 0
sim.n_cell = [64, 64, 1]
sim.blocking_factor_x = [16]
sim.blocking_factor_y = [16]
sim.blocking_factor_z = [1]
sim.prob_relative = [1.1]

# beam diagnostics
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = True

# domain decomposition & space charge mesh
sim.init_grids()

# beam energy ad bunch charge
kin_energy_MeV = 100  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 1000000  # number of macro particles

# intialize reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

# initialize particle bunch
distr = distribution.Gaussian(
    lambdaX=4.0e-5,
    lambdaY=4.0e-5,
    lambdaT=1.0e-3,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice)
ns = 1  # number of slices per ds in the element
L = 1.0  # drift length
sim.lattice.extend([monitor, elements.Drift(name="d1", ds=L, nslice=ns), monitor])

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Analyze

We run the following script to analyze correctness:

Listing 261 You can copy this file from examples/space_charge_gaussian_kick/analysis_sc_kick_Gauss2p5D.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#

import numpy as np
import openpmd_api as io
import scipy.constants as sc
from scipy.special import expi

# initial/final beam
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
last_step = list(series.iterations)[-1]
beam_initial = series.iterations[1].particles["beam"]
initial_sort = beam_initial.to_df().set_index("id")
beam_final = series.iterations[last_step].particles["beam"]
final_sort = beam_final.to_df().set_index("id")

# Physical constants
qe = sc.elementary_charge
me = sc.electron_mass
clite = sc.speed_of_light
eps0 = sc.epsilon_0
me_eV = me * clite**2 / qe

# Basic beam parameters
bunch_charge_C = 1.0e-9
kin_energy_MeV = 100
gamma = 1.0e6 * kin_energy_MeV / me_eV + 1.0
beta_gamma_2 = gamma**2 - 1.0
beta = np.sqrt(beta_gamma_2) / gamma

# Space charge intensity parameter (units of length, meters)
Kscale = (
    qe * bunch_charge_C / (4.0 * np.pi * eps0 * me * clite**2 * beta_gamma_2 * gamma)
)

# Beam distribution parameters
sigmax = 4.0e-5
sigmay = 4.0e-5
sigmact = 1.0e-3
sigmaz = beta * sigmact

# Simulation slice length
L = 1.0

# Longitudinal kick scale parameter
gauss_long_scale = 6.0

# Initial particle data

xi = initial_sort["position_x"]
pxi = initial_sort["momentum_x"]
yi = initial_sort["position_y"]
pyi = initial_sort["momentum_y"]
ti = initial_sort["position_t"]
pti = initial_sort["momentum_t"]
zi = -beta * ti

# Predicted momentum kick, from J. Qiang, Phys. Rev. Accel. Beams 28, 114602 (2025), eqs. (31-32)

ri_2 = xi**2 + yi**2
gauss_exp = np.exp(-ri_2 / (2.0 * sigmax**2))
gauss_xy_factor = 2.0 * (1.0 - gauss_exp) / ri_2
gauss_pdf_z = np.exp(-(zi**2) / (2.0 * sigmaz**2)) / (np.sqrt(2.0 * np.pi) * sigmaz)
px_predicted = L * Kscale * gauss_pdf_z * xi * gauss_xy_factor
py_predicted = L * Kscale * gauss_pdf_z * yi * gauss_xy_factor

potential_xy_factor = expi(-ri_2 / (2.0 * sigmax**2)) + np.log(
    gauss_long_scale**2 / ri_2
)
d_gauss_pdf_z = -zi / sigmaz**2 * gauss_pdf_z

pz_predicted = -L * Kscale * potential_xy_factor * d_gauss_pdf_z
pt_predicted = -beta * pz_predicted

# Maximum momentum kick
px_max = px_predicted.abs().max()
py_max = py_predicted.abs().max()
pt_max = pt_predicted.abs().max()

# Final particle data

xf = final_sort["position_x"]
pxf = final_sort["momentum_x"]
yf = final_sort["position_y"]
pyf = final_sort["momentum_y"]
tf = final_sort["position_t"]
ptf = final_sort["momentum_t"]

# Difference between value and prediction

dpx_max = (pxf - px_predicted).abs().max()
dpy_max = (pyf - py_predicted).abs().max()
dpt_max = (ptf - pt_predicted).abs().max()

dpx_rms = np.sqrt(np.mean(np.square(pxf - px_predicted)))
dpy_rms = np.sqrt(np.mean(np.square(pyf - py_predicted)))
dpt_rms = np.sqrt(np.mean(np.square(ptf - pt_predicted)))

print()
print("Difference between predicted and computed final momentum, absolute rms:")
print("dpx_rms", dpx_rms)
print("dpy_rms", dpy_rms)
print("dpt_rms", dpt_rms)

print()
print("Difference between predicted and computed final momentum, absolute max:")
print("dpx_max", dpx_max)
print("dpy_max", dpy_max)
print("dpt_max", dpt_max)

print()
print("Difference between predicted and computed final momentum (rms), relative:")
print("dpx_rms/px_max", dpx_rms / px_max)
print("dpy_rms/py_max", dpy_rms / py_max)
print("dpt_rms/pt_max", dpt_rms / pt_max)

# Test maximum error:
atol = 5.1e-2
print(f"  tol={atol}")

assert np.allclose(
    [dpx_rms / px_max, dpy_rms / py_max],
    [0.0, 0.0],
    atol=atol,
)

# Longitudinal kick is sensitive to noise (relax tolerance):
atol = 0.5
print(f"  tol={atol}")

assert np.allclose(
    [dpt_rms / pt_max],
    [0.0],
    atol=atol,
)


print()
print("Difference between predicted and computed final momentum (max), relative:")
print("dpx_max/px_max", dpx_max / px_max)
print("dpy_max/py_max", dpy_max / py_max)
print("dpt_max/pt_max", dpt_max / pt_max)

# Test maximum error:
atol = 5.1e-2
print(f"  tol={atol}")

assert np.allclose(
    [dpx_max / px_max, dpy_max / py_max],
    [0.0, 0.0],
    atol=atol,
)

# Longitudinal kick is sensitive to noise (relax tolerance):
atol = 1.1
print(f"  tol={atol}")

assert np.allclose(
    [dpt_max / pt_max],
    [0.0],
    atol=atol,
)
Listing 262 You can copy this file from examples/space_charge_gaussian_kick/analysis_sc_kick_PIC2p5D.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell, Ji Qiang
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = "2p5D"
sim.poisson_solver = "fft"
sim.space_charge_num_longitudinal_bins = 129
sim.space_charge_apply_longitudinal_kick = True
sim.slice_step_diagnostics = True

# set numerical parameters and IO control
sim.max_level = 0
sim.n_cell = [64, 64, 1]
sim.blocking_factor_x = [16]
sim.blocking_factor_y = [16]
sim.blocking_factor_z = [1]
sim.prob_relative = [1.1]

# beam diagnostics
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = True

# domain decomposition & space charge mesh
sim.init_grids()

# beam energy ad bunch charge
kin_energy_MeV = 100  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 1000000  # number of macro particles

# intialize reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

# initialize particle bunch
distr = distribution.Gaussian(
    lambdaX=4.0e-5,
    lambdaY=4.0e-5,
    lambdaT=1.0e-3,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice)
ns = 1  # number of slices per ds in the element
L = 1.0  # drift length
sim.lattice.extend([monitor, elements.Drift(name="d1", ds=L, nslice=ns), monitor])

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()