Expanding Beam in Free Space with 3D Space Charge

A coasting bunch expanding freely in free space under its own space charge.

We use a cold (zero emittance) 250 MeV electron bunch whose initial distribution is a uniformly-populated 3D ball of radius R0 = 1 mm when viewed in the bunch rest frame.

In the laboratory frame, the bunch is a uniformly-populated ellipsoid, which expands to twice its original size. This is tested using the second moments of the distribution.

In this test, the initial and final values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\epsilon_t\) must agree with nominal values.

This test uses mesh-refinement to solve the space charge force. The coarse grid wraps the beam maximum extent by 300%, emulating “open boundary” conditions. The refined grid in level 1 spans 110% of the beam maximum extent. The grid spacing is adaptively adjusted as the beam evolves.

Run

This example can be run either as:

  • Python script: python3 run_expanding_fft.py or

  • ImpactX executable using an input file: impactx input_expanding_fft.in

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

We also provide the same example with the multi-grid (MLMG) Poisson solver.

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

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.max_level = 1
sim.n_cell = [16, 16, 20]
sim.blocking_factor_x = [16]
sim.blocking_factor_y = [16]
sim.blocking_factor_z = [4]

sim.particle_shape = 2  # B-spline order
sim.space_charge = "3D"
sim.poisson_solver = "fft"
sim.dynamic_size = True
sim.prob_relative = [1.2, 1.1]

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

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

# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
kin_energy_MeV = 250  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles (outside tests, use 1e5 or more)

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

#   particle bunch
distr = distribution.Kurth6D(
    lambdaX=4.472135955e-4,
    lambdaY=4.472135955e-4,
    lambdaT=9.12241869e-7,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

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

# design the accelerator lattice
sim.lattice.extend([monitor, elements.Drift(name="d1", ds=6.0, nslice=40), monitor])

# run simulation
sim.track_particles()

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

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.max_level = 1
sim.n_cell = [16, 16, 20]
sim.blocking_factor_x = [16]
sim.blocking_factor_y = [16]
sim.blocking_factor_z = [4]

sim.particle_shape = 2  # B-spline order
sim.space_charge = "3D"
sim.poisson_solver = "multigrid"
sim.dynamic_size = True
sim.prob_relative = [3.0, 1.1]

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

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

# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
kin_energy_MeV = 250  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles (outside tests, use 1e5 or more)

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

#   particle bunch
distr = distribution.Kurth6D(
    lambdaX=4.472135955e-4,
    lambdaY=4.472135955e-4,
    lambdaT=9.12241869e-7,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

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

# design the accelerator lattice
sim.lattice.extend([monitor, elements.Drift(name="d1", ds=6.0, nslice=40), monitor])

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 214 You can copy this file from examples/expanding/input_expanding_fft.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000  # outside tests, use 1e5 or more
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = kurth6d
beam.lambdaX = 4.472135955e-4
beam.lambdaY = beam.lambdaX
beam.lambdaT = 9.12241869e-7
beam.lambdaPx = 0.0
beam.lambdaPy = 0.0
beam.lambdaPt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 monitor
lattice.nslice = 40

drift1.type = drift
drift1.ds = 6.0

monitor.type = beam_monitor
monitor.backend = h5


###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = 3D
algo.poisson_solver = "fft"

# Space charge solver with one MR level
amr.max_level = 1
amr.n_cell = 16 16 20
amr.blocking_factor_x = 16
amr.blocking_factor_y = 16
amr.blocking_factor_z = 4

geometry.prob_relative = 1.2 1.1

# Space charger solver without MR
#amr.max_level = 0
#amr.n_cell = 56 56 48
#geometry.prob_relative = 3.0
Listing 215 You can copy this file from examples/expanding/input_expanding_mlmg.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000  # outside tests, use 1e5 or more
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = kurth6d
beam.lambdaX = 4.472135955e-4
beam.lambdaY = beam.lambdaX
beam.lambdaT = 9.12241869e-7
beam.lambdaPx = 0.0
beam.lambdaPy = 0.0
beam.lambdaPt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 monitor
lattice.nslice = 40

drift1.type = drift
drift1.ds = 6.0

monitor.type = beam_monitor
monitor.backend = h5


###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = 3D
algo.poisson_solver = "multigrid"

# Space charge solver with one MR level
amr.max_level = 1
amr.n_cell = 16 16 20
amr.blocking_factor_x = 16
amr.blocking_factor_y = 16
amr.blocking_factor_z = 4

geometry.prob_relative = 3.0 1.1

# Space charger solver without MR
#amr.max_level = 0
#amr.n_cell = 56 56 48
#geometry.prob_relative = 3.0

Analyze

We run the following script to analyze correctness:

Script analysis_expanding.py
Listing 216 You can copy this file from examples/expanding/analysis_expanding.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
from scipy.stats import moment


def get_moments(beam):
    """Calculate standard deviations of beam position & momenta
    and emittance values

    Returns
    -------
    sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
    """
    sigx = moment(beam["position_x"], moment=2) ** 0.5  # variance -> std dev.
    sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
    sigy = moment(beam["position_y"], moment=2) ** 0.5
    sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
    sigt = moment(beam["position_t"], moment=2) ** 0.5
    sigpt = moment(beam["momentum_t"], moment=2) ** 0.5

    epstrms = beam.cov(ddof=0)
    emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
    emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
    emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5

    return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)


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

# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)

print("Initial Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)

atol = 0.0  # ignored
rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
    [
        4.4721359550e-004,
        4.4721359550e-004,
        9.1224186858e-007,
        0.0e-006,
        0.0e-006,
        0.0e-006,
    ],
    rtol=rtol,
    atol=atol,
)


print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)

atol = 0.0  # ignored
rtol = 1.6 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, sigt],
    [
        8.9442719100e-004,
        8.9442719100e-004,
        1.8244837370e-006,
    ],
    rtol=rtol,
    atol=atol,
)
atol = 1.0e-8
rtol = 0.0  # ignored
assert np.allclose(
    [emittance_x, emittance_y, emittance_t],
    [
        0.0,
        0.0,
        0.0,
    ],
    rtol=rtol,
    atol=atol,
)

Expanding Beam in Free Space with 2D Space Charge

A long, coasting unbunched beam expanding freely in free space under its own 2D space charge.

We use a cold (zero emittance) 250 MeV electron bunch whose initial distribution is a uniformly-populated cylinder of radius R0 = 1 mm.

In the laboratory frame, the beam expands to twice its original transverse size. This is tested using the second moments of the distribution.

In this test, the initial and final values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\epsilon_t\) must agree with nominal values.

Run

This example can be run either as:

  • Python script: python3 run_expanding_fft_2D.py or

  • ImpactX executable using an input file: impactx input_expanding_fft_2D.in

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

We also provide the same example with the multi-grid (MLMG) Poisson solver.

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

from impactx import ImpactX, distribution, elements

sim = ImpactX()

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

sim.particle_shape = 2  # B-spline order
sim.space_charge = "2D"
sim.poisson_solver = "fft"
sim.dynamic_size = True
sim.prob_relative = [1.1]

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

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

# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
kin_energy_MeV = 250  # reference energy
beam_current_A = 0.15  # beam current
npart = 10000  # number of macro particles (outside tests, use 1e5 or more)

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

#   particle bunch
distr = distribution.KVdist(
    lambdaX=5.0e-4,
    lambdaY=5.0e-4,
    lambdaT=1.0e-3,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
)
sim.add_particles(beam_current_A, distr, npart)

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

# design the accelerator lattice
doubling_distance = 10.612823669911099

sim.lattice.extend(
    [monitor, elements.Drift(name="d1", ds=doubling_distance, nslice=100), monitor]
)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 218 You can copy this file from examples/expanding/input_expanding_fft_2D.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000  # outside tests, use 1e5 or more
beam.units = static
beam.kin_energy = 250.0
beam.current = 0.15
beam.particle = proton
beam.distribution = kvdist
beam.lambdaX = 5.0e-4
beam.lambdaY = beam.lambdaX
beam.lambdaT = 1.0e-3  #result should not depend on this value
beam.lambdaPx = 0.0
beam.lambdaPy = 0.0
beam.lambdaPt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 monitor
lattice.nslice = 100

drift1.type = drift
drift1.ds = 10.612823669911099  #doubling-distance

monitor.type = beam_monitor
monitor.backend = h5


###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
#algo.track = "envelope"
algo.space_charge = 2D
algo.poisson_solver = "fft"

amr.n_cell = 32 32 1
amr.blocking_factor_x = 16
amr.blocking_factor_y = 16
amr.blocking_factor_z = 1

geometry.prob_relative = 1.1

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

We run the following script to analyze correctness:

Script analysis_expanding_fft_2D.py
Listing 219 You can copy this file from examples/expanding/analysis_expanding_fft_2D.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
from scipy.stats import moment


def get_moments(beam):
    """Calculate standard deviations of beam position & momenta
    and emittance values

    Returns
    -------
    sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
    """
    sigx = moment(beam["position_x"], moment=2) ** 0.5  # variance -> std dev.
    sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
    sigy = moment(beam["position_y"], moment=2) ** 0.5
    sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
    sigt = moment(beam["position_t"], moment=2) ** 0.5
    sigpt = moment(beam["momentum_t"], moment=2) ** 0.5

    epstrms = beam.cov(ddof=0)
    emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
    emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
    emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5

    return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)


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

# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)

print("Initial Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)

atol = 0.0  # ignored
rtol = 1.5 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
    [
        5.0e-004,
        5.0e-004,
        1.0e-003,
        0.0e-006,
        0.0e-006,
        0.0e-006,
    ],
    rtol=rtol,
    atol=atol,
)


print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)

atol = 0.0  # ignored
rtol = 3.5 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, sigt],
    [
        1.0e-003,
        1.0e-003,
        1.0e-003,
    ],
    rtol=rtol,
    atol=atol,
)
atol = 1.0e-8
rtol = 0.0  # ignored
assert np.allclose(
    [emittance_x, emittance_y, emittance_t],
    [
        0.0,
        0.0,
        0.0,
    ],
    rtol=rtol,
    atol=atol,
)

Expanding Beam in Free Space with 2D Space Charge, with Test Particles

A long, coasting unbunched beam expanding freely in free space under its own 2D space charge.

The problem parameters are identical to the above example, except that a set of test particles is loaded and tracked, together with the main particle bunch.

The test particle data is collected once per space charge slice, using the callback hook feature, and the test particle orbits are plotted vs. distance s.

Run

This example can be run as:

  • Python script: python3 run_expanding_fft_2D_test_particles.py or

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

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

import numpy as np
import pandas as pd

import amrex.space3d as amr
from impactx import ImpactX, distribution, elements

sim = ImpactX()

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

sim.particle_shape = 2  # B-spline order
sim.space_charge = "2D"
sim.poisson_solver = "fft"
sim.dynamic_size = True
sim.prob_relative = [1.1]

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

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

# basic beam properties
R0 = 1.0e-3  # initial beam radius (m)
Ib = 0.15  # beam current (A)
kin_energy_MeV = 250  # reference energy
beam_current_A = Ib  # beam current
npart = 10000  # number of macro particles (outside tests, use 1e5 or more)

#   reference particle
ref = sim.beam.ref
ref.set_species("proton").set_kin_energy_MeV(kin_energy_MeV)
qm_eev = ref.charge_qe / (ref.mass_MeV * 1.0e6)  # electron charge/mass in e / eV

#   particle bunch
distr = distribution.KVdist(
    lambdaX=R0 / 2.0,
    lambdaY=R0 / 2.0,
    lambdaT=1.0e-3,
    lambdaPx=0.0,
    lambdaPy=0.0,
    lambdaPt=0.0,
)
sim.add_particles(beam_current_A, distr, npart)

pc = sim.particle_container()

#  add test particles
if amr.ParallelDescriptor.IOProcessor():
    df_initial = pd.read_csv("./initial_coords.csv", sep=" ")
    dx, dy, dt, dpx, dpy, dpt = (
        df_initial[["x", "y", "t", "px", "py", "pt"]].to_numpy().T
    )
    pc.add_n_particles(dx, dy, dt, dpx, dpy, dpt, qm_eev, bunch_charge=0.0)

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

# design the accelerator lattice
betgam = ref.beta_gamma  # relativistic factor
IA = 3.1297388031196285e7  # alfven current for protons
Kpv = 2.0 * Ib / (IA * betgam**3)  # generalized beam perveance
double_constant = 1.516770632602484  # constant independent of beam parameters
doubling_distance = R0 * double_constant / np.sqrt(Kpv)
print("Doubling distance: ")
print(doubling_distance)

ns = 50  # recommended: for outside tests, using 100-200.

sim.lattice.extend(
    [monitor, elements.Drift(name="d1", ds=doubling_distance, nslice=ns), monitor]
)

sarr = []
test_data = []
mm_scale = 1.0e3


def hook_before_slice(sim):
    s = sim.beam.ref.s
    sarr.append(s)
    beam = sim.beam.to_df()
    # Filter on particle weight (collect test particles only)
    for row in beam[beam["weighting"] == 0.0].itertuples():
        # collect test particle data
        test_data.append(
            [s, row.idcpu, row.position_x * mm_scale, row.position_y * mm_scale]
        )


sim.hook["before_slice"] = hook_before_slice

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

df = pd.DataFrame(test_data, columns=["s", "id", "x", "y"])
sorted_df = df.sort_values(by="id")

# Uncomment the lines below to generate a plot of the test particle orbits:
# n = len(sarr)
# for i in range(0, len(df), n):
#    subset = sorted_df.iloc[i : i + n]
#    plt.scatter(subset["s"], subset["x"], s=5)

# plt.xlabel("s [m]", fontsize=12)
# plt.ylabel("x [mm]", fontsize=12)
do_plot = False  # True to generate a plot of the test particle orbits
if do_plot:
    import matplotlib.pyplot as plt

    n = len(sarr)
    for i in range(0, len(df), n):
        subset = sorted_df.iloc[i : i + n]
        plt.scatter(subset["s"], subset["x"], s=5)

    plt.xlabel("s [m]", fontsize=12)
    plt.ylabel("x [mm]", fontsize=12)
    plt.title("Test Particles: Horizontal Coordinates")
# plt.show()

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