Symplectic Integration in an Exact Quadrupole

This benchmark tests the use of the ExactQuad (quadrupole_exact) element for integrating through a quadrupole using the exact nonlinear Hamiltonian.

A 25 pC electron bunch with 100 MeV total energy, a small initial rms beam size of (3.9, 3.9, 1.0) microns, 2 mrad transverse divergence and 2.5% energy spread undergoes rapid expansion followed by transverse focusing in a quadrupole doublet. The parameters are chosen such that chromatic focusing effects are important.

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.

In addition, the Hamiltonian value is computed for each particle at the entrance and exit of the final quadrupole. The change in the Hamiltonian value, taken relative to the standard deviation \(\sigma_H\) over particles, must be smaller than the allowed tolerance (here, taken to be 0.1%).

Run

This example can be run either as:

  • Python script: python3 run_exact_quad.py or

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

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

Listing 121 You can copy this file from examples/symplectic_integration/run_exact_quad.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
from scipy.constants import c, e, m_e

from impactx import ImpactX, distribution, elements, twiss

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
sim.slice_step_diagnostics = True

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

# basic beam parameters
total_energy_MeV = 100.0  # reference energy (total)
mass_MeV = 0.510998950  # particle mass
kin_energy_MeV = total_energy_MeV - mass_MeV
bunch_charge_C = 25.0e-12  # used with space charge
npart = 10000  # number of macro particles

# set reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(-1.0).set_mass_MeV(mass_MeV).set_kin_energy_MeV(kin_energy_MeV)

# factors converting the beam distribution to ImpactX input
gamma = total_energy_MeV / mass_MeV
bg = np.sqrt(gamma**2 - 1.0)
sigma_tau = 1e-6  # in m
sigma_p = 2.5e-2  # dimensionless
rigidity = m_e * c * bg / e

#   particle bunch
distr = distribution.Triangle(
    **twiss(
        beta_x=0.002,
        beta_y=0.002,
        beta_t=sigma_tau / sigma_p,
        emitt_x=1.5e-6 / bg,
        emitt_y=1.5e-6 / bg,
        emitt_t=sigma_tau * sigma_p,
        alpha_x=0.0,
        alpha_y=0.0,
        alpha_t=0.0,
    )
)
sim.add_particles(bunch_charge_C, distr, npart)

# design the accelerator lattice
ns = 1  # number of slices per ds in the element

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

# element names consistent with HTU_base_lattice_matched.lte Elegant input

drift1 = elements.ExactDrift(name="drift1", ds=0.046, nslice=ns)
quad1 = elements.ExactQuad(
    name="quad1", ds=0.02903, k=207.0, unit=1, mapsteps=5, nslice=ns
)
quad2 = elements.ExactQuad(
    name="quad2", ds=0.02890, k=-207.0, unit=1, mapsteps=5, nslice=ns
)

# set the lattice
sim.lattice.append(monitor0)
sim.lattice.append(drift1)
sim.lattice.append(quad1)
sim.lattice.append(monitor1)
sim.lattice.append(quad2)
sim.lattice.append(monitor1)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 122 You can copy this file from examples/symplectic_integration/input_exact_quad.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 99.48900104
beam.charge = 25.0e-12
beam.particle = electron
beam.distribution = triangle_from_twiss
beam.alphaX = 0.0
beam.alphaY = 0.0
beam.alphaT = 0.0
beam.betaX = 0.002
beam.betaY = beam.betaX
beam.betaT = 0.00004
beam.emittX = 7.626114e-9
beam.emittY = beam.emittX
beam.emittT = 2.5e-08


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor0 drift1 quad1 monitor1 quad2 monitor1
lattice.nslice = 1

monitor0.type = beam_monitor
monitor0.name = "monitor0"
monitor0.backend = h5

monitor1.type = beam_monitor
monitor1.name = "monitor"
monitor1.backend = h5

drift1.type = drift_exact
drift1.ds = 0.046

quad1.type = quad_exact
quad1.ds = 0.02903
quad1.k = 207.0
quad1.units = 1
quad1.mapsteps = 5

quad2.type = quad_exact
quad2.ds = 0.02903
quad2.k = -207.0
quad2.units = 1
quad2.mapsteps = 5

###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false


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

Analyze

We run the following script to analyze correctness:

Script analysis_exact_quad.py
Listing 123 You can copy this file from examples/symplectic_integration/analysis_exact_quad.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.constants import c, e, m_e
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)

initial_step = list(series.iterations)[0]
last_step = list(series.iterations)[-1]
initial = series.iterations[initial_step].particles["beam"].to_df()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.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 = 3.0 * 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],
    [
        1.768569e-04,
        1.195967e-04,
        1.034426e-06,
        1.356459e-08,
        9.374827e-09,
        2.581087e-08,
    ],
    rtol=rtol,
    atol=atol,
)

print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
s_ref = beam_final.get_attribute("s_ref")
gamma_ref = beam_final.get_attribute("gamma_ref")
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}\n"
    f"  s_ref={s_ref:e} gamma_ref={gamma_ref:e}"
)

atol = 0.0  # ignored
rtol = 2.2 * 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],
    [
        2.458574e-04,
        1.513417e-04,
        1.068497e-06,
        1.797301e-08,
        1.009999e-08,
        2.661524e-08,
    ],
    rtol=rtol,
    atol=atol,
)


# join tables on particle ID, so we can compare the same particle initial->final
beam_joined = final.join(initial, lsuffix="_final", rsuffix="_initial")
xi = beam_joined["position_x_initial"]
yi = beam_joined["position_y_initial"]
pxi = beam_joined["momentum_x_initial"]
pyi = beam_joined["momentum_y_initial"]
pti = beam_joined["momentum_t_initial"]
xf = beam_joined["position_x_final"]
yf = beam_joined["position_y_final"]
pxf = beam_joined["momentum_x_final"]
pyf = beam_joined["momentum_y_final"]
ptf = beam_joined["momentum_t_final"]

# Parameters appearing in the Hamiltonian
beta_ref = beam_final.get_attribute("beta_ref")
bg_ref = beam_final.get_attribute("beta_gamma_ref")
rigidity = m_e * bg_ref * c / e
B_gradient = 207.0
k = B_gradient / rigidity

# Evaluate the change in Hamiltonian value for each particle
Hi = (
    -np.sqrt(1.0 - 2.0 / beta_ref * pti + pti**2 - pxi**2 - pyi**2)
    + (k / 2.0) * (xi**2 - yi**2)
    - pti / beta_ref
    + 1.0
)
Hf = (
    -np.sqrt(1.0 - 2.0 / beta_ref * ptf + ptf**2 - pxf**2 - pyf**2)
    + (k / 2.0) * (xf**2 - yf**2)
    - ptf / beta_ref
    + 1.0
)

H_sigma = (np.std(Hi) + np.std(Hf)) / 2.0
dH = (Hf - Hi).abs()
dH_max_relative = dH.max() / H_sigma

# particle-wise comparison of H & I initial to final
atol = 1.0e-3
rtol = 0.0  # large number
print()
print(f"  atol={atol} (ignored: rtol~={rtol})")

print(f"  Standard deviation in H: H_sigma = {H_sigma}")
print(f"  dH_max relative to H_sigma = {dH_max_relative}")
assert np.allclose(dH_max_relative, 0.0, rtol=rtol, atol=atol)

FODO Channel with Quads Treated as Exact Multipoles

This is identical to examples-fodo, except that the quadrupoles have been replaced by ExactMultipole (multipole_exact) elements with only a nonzero quadrupole coefficient. Its purpose is to test the limiting case of the ExactMultipole model corresponding to a simple quadrupole.

The analysis script is identical to the analysis script used for examples-fodo.

Run

This example can be run either as:

  • Python script: python3 run_fodo_multipole.py or

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

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

Listing 124 You can copy this file from examples/symplectic_integration/run_fodo_multipole.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.space_charge = False
# 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 = 2.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=3.9984884770e-5,
    lambdaY=3.9984884770e-5,
    lambdaT=1.0e-3,
    lambdaPx=2.6623538760e-5,
    lambdaPy=2.6623538760e-5,
    lambdaPt=2.0e-3,
    muxpx=-0.846574929020762,
    muypy=0.846574929020762,
    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 = 25  # number of slices per ds in the element
fodo = [
    monitor,
    elements.Drift(name="drift1", ds=0.25, nslice=ns),
    monitor,
    elements.ExactMultipole(
        name="quad1",
        ds=1.0,
        k_normal=[0.0, 1.0],
        k_skew=[0.0, 0.0],
        unit=0,
        mapsteps=5,
        nslice=ns,
    ),
    monitor,
    elements.Drift(name="drift2", ds=0.5, nslice=ns),
    monitor,
    elements.ExactMultipole(
        name="quad2",
        ds=1.0,
        k_normal=[0.0, -1.0],
        k_skew=[0.0, 0.0],
        unit=0,
        mapsteps=5,
        nslice=ns,
    ),
    monitor,
    elements.Drift(name="drift3", ds=0.25, nslice=ns),
    monitor,
]
# assign a fodo segment
sim.lattice.extend(fodo)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 125 You can copy this file from examples/symplectic_integration/input_fodo_multipole.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 3.9984884770e-5
beam.lambdaY = beam.lambdaX
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-5
beam.lambdaPy = beam.lambdaPx
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = -beam.muxpx
beam.mutpt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift3 monitor
lattice.nslice = 10

monitor.type = beam_monitor
monitor.backend = h5

drift1.type = drift
drift1.ds = 0.25

quad1.type = multipole_exact
quad1.ds = 1.0
quad1.k_normal = 0 1.0
quad1.k_skew = 0 0
quad1.units = 0
quad1.int_order = 4
quad1.mapsteps = 4

drift2.type = drift
drift2.ds = 0.5

quad2.type = multipole_exact
quad2.ds = 1.0
quad2.k_normal = 0 -1.0
quad2.k_skew = 0 0
quad2.units = 0
quad2.int_order = 4
quad2.mapsteps = 4

drift3.type = drift
drift3.ds = 0.25


###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
#algo.track = envelope  #to test in envelope mode

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

Analyze

We run the following script to analyze correctness:

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

import glob
import re

import numpy as np
import pandas as pd


def read_file(file_pattern):
    for filename in glob.glob(file_pattern):
        df = pd.read_csv(filename, delimiter=r"\s+")
        if "step" not in df.columns:
            step = int(re.findall(r"[0-9]+", filename)[0])
            df["step"] = step
        yield df


def read_time_series(file_pattern):
    """Read in all CSV files from each MPI rank (and potentially OpenMP
    thread). Concatenate into one Pandas dataframe.

    Returns
    -------
    pandas.DataFrame
    """
    return pd.concat(
        read_file(file_pattern),
        axis=0,
        ignore_index=True,
    )  # .set_index('id')


# read reduced diagnostics
rbc = read_time_series("diags/reduced_beam_characteristics.*")

s = rbc["s"]
sigma_x = rbc["sigma_x"]
sigma_y = rbc["sigma_y"]
sigma_t = rbc["sigma_t"]
emittance_x = rbc["emittance_x"]
emittance_y = rbc["emittance_y"]
emittance_t = rbc["emittance_t"]

sig_xi = sigma_x.iloc[0]
sig_yi = sigma_y.iloc[0]
sig_ti = sigma_t.iloc[0]
emittance_xi = emittance_x.iloc[0]
emittance_yi = emittance_y.iloc[0]
emittance_ti = emittance_t.iloc[0]

length = len(s) - 1

sf = s.iloc[length]
sig_xf = sigma_x.iloc[length]
sig_yf = sigma_y.iloc[length]
sig_tf = sigma_t.iloc[length]
emittance_xf = emittance_x.iloc[length]
emittance_yf = emittance_y.iloc[length]
emittance_tf = emittance_t.iloc[length]


print("Initial Beam:")
print(f"  sigx={sig_xi:e} sigy={sig_yi:e} sigt={sig_ti:e}")
print(
    f"  emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}"
)

atol = 0.0  # ignored
rtol = 1.0e-2  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti],
    [
        7.5451170454175073e-005,
        7.5441588239210947e-005,
        9.9775878164077539e-004,
        1.9959540393751392e-009,
        2.0175015289132990e-009,
        2.0013820193294972e-006,
    ],
    rtol=rtol,
    atol=atol,
)


print("")
print("Final Beam:")
print(f"  sigx={sig_xf:e} sigy={sig_yf:e} sigt={sig_tf:e}")
print(
    f"  emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}"
)

atol = 0.0  # ignored
rtol = 1.0e-2  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf],
    [
        7.4790118496224206e-005,
        7.5357525169680140e-005,
        9.9775879288128088e-004,
        1.9959539836392703e-009,
        2.0175014668882125e-009,
        2.0013820380883801e-006,
    ],
    rtol=rtol,
    atol=atol,
)

Symplectic Integration in a Long Sextupole

This benchmark tests the use of the ExactMultipole (multipole_exact) element for integrating through a long sextupole using the exact nonlinear Hamiltonian.

An array of initial conditions corresponding to protons with 0.8 GeV total energy is tracked through a 0.5 m long sextupole.

In this test, each particle’s final phase space vector is compared against numerical tracking results obtained in PTC (E. Stern).

All 6 phase space coordinates must agree within the specified tolerance.

Run

This example can be run as:

  • Python script: python3 run_sextupole.py

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

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

import pandas as pd

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

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
sim.slice_step_diagnostics = True

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

# load initial beam parameters
kin_energy_MeV = 0.8e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
qm_eev = 1.0 / 938.27208816 / 1e6  # electron charge/mass in e / eV

pc = sim.particle_container()

df_initial = pd.read_csv("./initial_coords.csv", sep=" ")
dx = df_initial["x"].to_numpy()
dpx = df_initial["px"].to_numpy()
dy = df_initial["y"].to_numpy()
dpy = df_initial["py"].to_numpy()
dt = df_initial["t"].to_numpy()
dpt = df_initial["pt"].to_numpy()
if not Config.have_gpu:  # initialize using cpu-based PODVectors
    dx_podv = amr.PODVector_real_std()
    dy_podv = amr.PODVector_real_std()
    dt_podv = amr.PODVector_real_std()
    dpx_podv = amr.PODVector_real_std()
    dpy_podv = amr.PODVector_real_std()
    dpt_podv = amr.PODVector_real_std()
else:  # initialize on device using arena/gpu-based PODVectors
    dx_podv = amr.PODVector_real_arena()
    dy_podv = amr.PODVector_real_arena()
    dt_podv = amr.PODVector_real_arena()
    dpx_podv = amr.PODVector_real_arena()
    dpy_podv = amr.PODVector_real_arena()
    dpt_podv = amr.PODVector_real_arena()

for p_dx in dx:
    dx_podv.push_back(p_dx)
for p_dy in dy:
    dy_podv.push_back(p_dy)
for p_dt in dt:
    dt_podv.push_back(p_dt)
for p_dpx in dpx:
    dpx_podv.push_back(p_dpx)
for p_dpy in dpy:
    dpy_podv.push_back(p_dpy)
for p_dpt in dpt:
    dpt_podv.push_back(p_dpt)

pc.add_n_particles(
    dx_podv, dy_podv, dt_podv, dpx_podv, dpy_podv, dpt_podv, qm_eev, bunch_charge_C
)

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

# design the accelerator lattice)
ns = 1  # number of slices per ds in the element
line = [
    monitor,
    elements.ExactMultipole(
        name="sextupole1",
        ds=0.5,
        k_normal=[0.0, 0.0, 0.35],
        k_skew=[0.0, 0.0, 0.0],
        unit=0,
        int_order=4,
        mapsteps=5,
        nslice=ns,
    ),
    monitor,
]
# assign a fodo segment
sim.lattice.extend(line)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Analyze

We run the following script to analyze correctness:

Script analysis_sextupole.py
Listing 128 You can copy this file from examples/symplectic_integration/analysis_sextupole.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 pandas as pd

# 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()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.to_df()

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

# load particle data
df_initial = pd.read_csv("./initial_coords.csv", sep=" ")
df_final = pd.read_csv("./final_coords.csv", sep=" ")

# compute differences
error_xi = (df_initial["x"] - initial["position_x"]).abs().max()
error_pxi = (df_initial["px"] - initial["momentum_x"]).abs().max()
error_yi = (df_initial["y"] - initial["position_y"]).abs().max()
error_pyi = (df_initial["py"] - initial["momentum_y"]).abs().max()
error_ti = (df_initial["t"] - initial["position_t"]).abs().max()
error_pti = (df_initial["pt"] - initial["momentum_t"]).abs().max()

error_xf = (df_final["x"] - final["position_x"]).abs().max()
error_pxf = (df_final["px"] - final["momentum_x"]).abs().max()
error_yf = (df_final["y"] - final["position_y"]).abs().max()
error_pyf = (df_final["py"] - final["momentum_y"]).abs().max()
error_tf = (df_final["t"] - final["position_t"]).abs().max()
error_ptf = (df_final["pt"] - final["momentum_t"]).abs().max()

xf_max = df_final["x"].abs().max()
pxf_max = df_final["px"].abs().max()
yf_max = df_final["y"].abs().max()
pyf_max = df_final["py"].abs().max()
tf_max = df_final["t"].abs().max()
ptf_max = df_final["pt"].abs().max()

print("Initial beam, maximum absolute difference in each coordinate:")
print("Difference x, px, y, py, t, pt:")
print(error_xi)
print(error_pxi)
print(error_yi)
print(error_pyi)
print(error_ti)
print(error_pti)

atol = 1.0e-13
print(f"  atol={atol}")

assert np.allclose(
    [error_xi, error_pxi, error_yi, error_pyi, error_ti, error_pti],
    [
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
    ],
    atol=atol,
)

print("")
print("Final beam, maximum relative difference in transverse coordinates:")
print("Difference x, px, y, py:")
print(error_xf / xf_max)
print(error_pxf / pxf_max)
print(error_yf / yf_max)
print(error_pyf / pyf_max)

atol = 1.0e-7
print(f"  tol={atol}")

assert np.allclose(
    [error_xf / xf_max, error_pxf / pxf_max, error_yf / yf_max, error_pyf / pyf_max],
    [
        0.0,
        0.0,
        0.0,
        0.0,
    ],
    atol=atol,
)

print("")
print("Final beam, maximum absolute difference in longitudinal coordinates:")
print("Difference t, pt:")
print(error_tf)
print(error_ptf)

atol = 1.0e-13
print(f"  atol={atol}")

assert np.allclose(
    [error_tf, error_ptf],
    [
        0.0,
        0.0,
    ],
    atol=atol,
)

Symplectic Integration in an Exact Combined-Function Bend

This benchmark tests the use of the ExactCFbend (cfbend_exact) element for integrating through a combined-function bend using the exact nonlinear Hamiltonian.

This example tests the transport of a 2 GeV electron bunch through a combined function bending element with only the lowest-order (dipole) coefficient nonzero, representing the effect of a pure dipole field. This is compared with the effect of the ExactSbend element for equivalent field strength, by applying the inverse map (ds < 0 and phi < 0).

As a result, the second moments of x, y, and t and the associated emittances of the bunch (as well as individual particle coordinates) should all be exactly unchanged.

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, to within the specified tolerance.

Run

This example can be run either as:

  • Python script: python3 run_exact_cfbend.py or

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

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

Listing 129 You can copy this file from examples/symplectic_integration/run_exact_cfbend.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.particle_shape = 2  # B-spline order
sim.space_charge = False
sim.slice_step_diagnostics = True

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

# basic beam parameters
mass_MeV = 0.510998950  # particle mass
kin_energy_MeV = 2.0e3
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

# set reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(-1.0).set_mass_MeV(mass_MeV).set_kin_energy_MeV(kin_energy_MeV)

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

# design the accelerator lattice
ns = 1  # number of slices per ds in the element

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

# lattice elements
cfbend1 = elements.ExactCFbend(
    name="cfbend1",
    ds=1.0,
    k_normal=[0.1, 0.0],
    k_skew=[0.0, 0.0],
    unit=0,
    int_order=2,
    mapsteps=5,
    nslice=ns,
)
sbend1 = elements.ExactSbend(name="sbend1", ds=-1.0, phi=-5.729577951308232, nslice=ns)

# set the lattice
sim.lattice.append(monitor)
sim.lattice.append(cfbend1)
sim.lattice.append(sbend1)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 130 You can copy this file from examples/symplectic_integration/input_exact_cfbend.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 4.0e-5
beam.lambdaY = beam.lambdaX
beam.lambdaT = 1.0e-3
beam.lambdaPx = 3.0e-5
beam.lambdaPy = beam.lambdaPx
beam.lambdaPt = 2.0e-4
beam.muxpx = 0.0
beam.muypy = -beam.muxpx
beam.mutpt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor cfbend1 sbend1 monitor
lattice.nslice = 10

monitor.type = beam_monitor
monitor.backend = h5

cfbend1.type = cfbend_exact
cfbend1.ds = 1.0
cfbend1.k_normal = 0.1 0.0
cfbend1.k_skew = 0 0
cfbend1.units = 0
cfbend1.int_order = 2
cfbend1.mapsteps = 5

sbend1.type = sbend_exact
sbend1.ds = -1.0
sbend1.phi = -5.729577951308232

###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false