Dogleg

This is a 2-bend dogleg lattice obtained by taking the first 1/2 of the Berlin-Zeuthen magnetic bunch compression chicane: https://www.desy.de/csr/

The primary purpose is to benchmark the reduced beam diagnostics in lattice regions with nonzero dispersion.

All parameters can be found online. A 5 GeV electron bunch with normalized transverse rms emittance of 1 um is used.

The final expected dispersion is 267 mm.

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 initial and final values of \(\alpha_x\), \(\alpha_y\), \(\beta_x\), \(\beta_y\), \(D_x\), and \(D_{px}\) must agree with nominal values.

Run

This example can be run either as:

• Python script: python3 run_dogleg.py or

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

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

Python: Script

Listing 117 You can copy this file from examples/dogleg/run_dogleg.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Marco Garten, 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 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 5.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

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

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=2.2951017632e-5,
    lambdaY=1.3084093142e-5,
    lambdaT=5.5555553e-8,
    lambdaPx=1.598353425e-6,
    lambdaPy=2.803697378e-6,
    lambdaPt=2.000000000e-6,
    muxpx=0.933345606203060,
    muypy=0.933345606203060,
    mutpt=0.999999961419755,
)
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
rc = 10.3462283686195526  # bend radius (meters)
psi = 0.048345620280243  # pole face rotation angle (radians)
lb = 0.500194828041958  # bend arc length (meters)

# Drift elements
dr1 = elements.Drift(name="dr1", ds=5.0058489435, nslice=ns)
dr2 = elements.Drift(name="dr2", ds=0.5, nslice=ns)

# Bend elements
sbend1 = elements.Sbend(name="sbend1", ds=lb, rc=-rc, nslice=ns)
sbend2 = elements.Sbend(name="sbend2", ds=lb, rc=rc, nslice=ns)

# Dipole Edge Focusing elements
dipedge1 = elements.DipEdge(name="dipedge1", psi=-psi, rc=-rc, g=0.0, K2=0.0)
dipedge2 = elements.DipEdge(name="dipedge2", psi=psi, rc=rc, g=0.0, K2=0.0)

lattice_dogleg = [sbend1, dipedge1, dr1, dipedge2, sbend2, dr2]

sim.lattice.append(monitor)
sim.lattice.extend(lattice_dogleg)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Executable: Input File

Listing 118 You can copy this file from examples/dogleg/input_dogleg.in.

###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 5.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 2.2951017632e-5
beam.lambdaY = 1.3084093142e-5
beam.lambdaT = 5.5555553e-8
beam.lambdaPx = 1.598353425e-6
beam.lambdaPy = 2.803697378e-6
beam.lambdaPt = 2.000000000e-6
beam.muxpx = 0.933345606203060
beam.muypy = beam.muxpx
beam.mutpt = 0.999999961419755


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 monitor
lattice.nslice = 25

sbend1.type = sbend
sbend1.ds = 0.500194828041958       # projected length 0.5 m, angle 2.77 deg
sbend1.rc = -10.3462283686195526

drift1.type = drift
drift1.ds = 5.0058489435            # projected length 5 m

sbend2.type = sbend
sbend2.ds = sbend1.ds               # projected length 0.5 m, angle 2.77 deg
sbend2.rc = -sbend1.rc

drift2.type = drift
drift2.ds = 0.5

dipedge1.type = dipedge   # dipole edge focusing
dipedge1.psi = -0.048345620280243
dipedge1.rc = -10.3462283686195526
dipedge1.g = 0.0
dipedge1.K2 = 0.0

dipedge2.type = dipedge
dipedge2.psi = -dipedge1.psi
dipedge2.rc = -dipedge1.rc
dipedge2.g = 0.0
dipedge2.K2 = 0.0

monitor.type = beam_monitor
monitor.backend = h5


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


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

Analyze

We run the following script to analyze correctness:

Script analysis_dogleg.py

Listing 119 You can copy this file from examples/dogleg/analysis_dogleg.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)


def get_twiss(openpmd_beam):
    """Return Twiss functions from an openPMD particle species

    Returns
    -------
    alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px
    """
    alpha_x = openpmd_beam.get_attribute("alpha_x")
    beta_x = openpmd_beam.get_attribute("beta_x")
    d_x = openpmd_beam.get_attribute("dispersion_x")
    d_px = openpmd_beam.get_attribute("dispersion_px")
    alpha_y = openpmd_beam.get_attribute("alpha_y")
    beta_y = openpmd_beam.get_attribute("beta_y")
    # d_y = openpmd_beam.get_attribute("dispersion_y")
    # d_py = openpmd_beam.get_attribute(["dispersion_py")

    return (alpha_x, beta_x, alpha_y, beta_y, d_x, d_px)


# 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_beam = series.iterations[last_step].particles["beam"]
final = final_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 = 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],
    [
        6.4214719960819659e-005,
        3.6603372435649773e-005,
        1.9955175623579313e-004,
        1.0198263116327677e-010,
        1.0308359092878036e-010,
        4.0035161705244885e-010,
    ],
    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 = 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],
    [
        1.922660e-03,
        2.166654e-05,
        1.101353e-04,
        8.561046e-09,
        1.020439e-10,
        8.569865e-09,
    ],
    rtol=rtol,
    atol=atol,
)

print("")
print("Final Twiss functions:")
alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px = get_twiss(final_beam)
print(f"  alpha_x={alpha_x:e} beta_x={beta_x:e} alpha_y={alpha_y:e} beta_y={beta_y:e}")
print(f"  dispersion_x={dispersion_x:e} dispersion_px={dispersion_px: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(
    [alpha_x, beta_x, alpha_y, beta_y, dispersion_x],
    [
        1.338583e00,
        1.437407e01,
        -1.347910e00,
        4.600362e00,
        -2.666972e-01,
    ],
    rtol=rtol,
    atol=atol,
)

Dogleg in Reverse

This is the reverse of the 2-bend dogleg lattice, obtained by taking the second 1/2 of the Berlin-Zeuthen magnetic bunch compression chicane: https://www.desy.de/csr/

The primary purpose is to demonstrate the initialization of a beam with a nonzero x-pt correlation in a lattice region with nonzero dispersion.

In this example, the x-pt correlation of the beam is removed at the dogleg exit, and the dispersion goes to zero.

The initial dispersion is taken to be -267 mm.

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 initial and final values of \(\alpha_x\), \(\alpha_y\), \(\beta_x\), \(\beta_y\), \(D_x\), and \(D_{px}\) must agree with nominal values.

Run

This example can be run either as:

• Python script: python3 run_dogleg_reverse.py or

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

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

Python: Script

Listing 120 You can copy this file from examples/dogleg/run_dogleg_reverse.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Marco Garten, Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements, twiss

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 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 5.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

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

#   particle bunch
distr = distribution.Waterbag(
    **twiss(
        beta_x=14.374073092185883,
        beta_y=4.6003616743069093,
        beta_t=1.4153999755977573,
        emitt_x=8.180911194584674375e-11,
        emitt_y=1.020438927769593e-10,
        emitt_t=8.5698645128105327e-09,
        alpha_x=1.3385830279518021,
        alpha_y=-1.3479109197361046,
        alpha_t=92.624347459169869,
        dispersion_x=-0.26669723385388505,
    ),
)
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
rc = 10.3462283686195526  # bend radius (meters)
psi = 0.048345620280243  # pole face rotation angle (radians)
lb = 0.500194828041958  # bend arc length (meters)

# Drift elements
dr1 = elements.Drift(name="dr1", ds=5.0058489435, nslice=ns)
dr2 = elements.Drift(name="dr2", ds=0.5, nslice=ns)

# Bend elements
sbend1 = elements.Sbend(name="sbend1", ds=lb, rc=-rc, nslice=ns)
sbend2 = elements.Sbend(name="sbend2", ds=lb, rc=rc, nslice=ns)

# Dipole Edge Focusing elements
dipedge1 = elements.DipEdge(name="dipedge1", psi=-psi, rc=-rc, g=0.0, K2=0.0)
dipedge2 = elements.DipEdge(name="dipedge2", psi=psi, rc=rc, g=0.0, K2=0.0)

lattice_dogleg = [sbend1, dipedge1, dr1, dipedge2, sbend2, dr2]

sim.lattice.append(monitor)
lattice_dogleg.reverse()
sim.lattice.extend(lattice_dogleg)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Executable: Input File

Listing 121 You can copy this file from examples/dogleg/input_dogleg_reverse.in.

###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 5.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag_from_twiss
beam.alphaX = 1.33858302795180
beam.alphaY = -1.347910457923642
beam.alphaT = 92.624347459105241
beam.betaX = 14.37407309218588
beam.betaY = 4.600362059144475
beam.betaT = 1.4153999755967554
beam.emittX = 8.180911194584674375e-11
beam.emittY = 1.020438927769593e-10
beam.emittT = 8.5698645128164239e-09
beam.dispX = -0.2667

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 monitor
lattice.reverse = true
lattice.nslice = 25

sbend1.type = sbend
sbend1.ds = 0.500194828041958       # projected length 0.5 m, angle 2.77 deg
sbend1.rc = -10.3462283686195526

drift1.type = drift
drift1.ds = 5.0058489435            # projected length 5 m

sbend2.type = sbend
sbend2.ds = sbend1.ds               # projected length 0.5 m, angle 2.77 deg
sbend2.rc = -sbend1.rc

drift2.type = drift
drift2.ds = 0.5

dipedge1.type = dipedge   # dipole edge focusing
dipedge1.psi = -0.048345620280243
dipedge1.rc = -10.3462283686195526
dipedge1.g = 0.0
dipedge1.K2 = 0.0

dipedge2.type = dipedge
dipedge2.psi = -dipedge1.psi
dipedge2.rc = -dipedge1.rc
dipedge2.g = 0.0
dipedge2.K2 = 0.0

monitor.type = beam_monitor
monitor.backend = h5


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


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

Analyze

We run the following script to analyze correctness:

Script analysis_dogleg_reverse.py

Listing 122 You can copy this file from examples/dogleg/analysis_dogleg_reverse.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)


def get_twiss(openpmd_beam):
    """Return Twiss functions from an openPMD particle species

    Returns
    -------
    alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px
    """
    alpha_x = openpmd_beam.get_attribute("alpha_x")
    beta_x = openpmd_beam.get_attribute("beta_x")
    d_x = openpmd_beam.get_attribute("dispersion_x")
    d_px = openpmd_beam.get_attribute("dispersion_px")
    alpha_y = openpmd_beam.get_attribute("alpha_y")
    beta_y = openpmd_beam.get_attribute("beta_y")
    # d_y = openpmd_beam.get_attribute("dispersion_y")
    # d_py = openpmd_beam.get_attribute(["dispersion_py")

    return (alpha_x, beta_x, alpha_y, beta_y, d_x, d_px)


# initial/final beam
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
last_step = list(series.iterations)[-1]
initial_beam = series.iterations[1].particles["beam"]
initial = initial_beam.to_df()
final_beam = series.iterations[last_step].particles["beam"]
final = final_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 = 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],
    [
        1.924711e-03,
        2.165646e-05,
        1.102534e-04,
        7.668809e-09,
        1.018986e-10,
        8.588054e-09,
    ],
    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 = 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.057123e-05,
        6.911405e-05,
        2.012573e-05,
        8.174618e-11,
        1.018986e-10,
        1.151058e-08,
    ],
    rtol=rtol,
    atol=atol,
)


print("")
print("Initial Twiss functions:")
alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px = get_twiss(initial_beam)
print(f"  alpha_x={alpha_x:e} beta_x={beta_x:e} alpha_y={alpha_y:e} beta_y={beta_y:e}")
print(f"  dispersion_x={dispersion_x:e} dispersion_px={dispersion_px: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(
    [alpha_x, beta_x, alpha_y, beta_y, dispersion_x],
    [
        1.340770e00,
        1.440253e01,
        -1.347747e00,
        4.602637e00,
        -2.667038e-01,
    ],
    rtol=rtol,
    atol=atol,
)

print("")
print("Final Twiss functions:")
alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px = get_twiss(final_beam)
print(f"  alpha_x={alpha_x:e} beta_x={beta_x:e} alpha_y={alpha_y:e} beta_y={beta_y:e}")
print(f"  dispersion_x={dispersion_x:e} dispersion_px={dispersion_px: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(
    [beta_x, alpha_y, beta_y],
    [
        5.176642e00,
        -4.973010e00,
        4.687750e01,
    ],
    rtol=rtol,
    atol=atol,
)

# We use absolute tolerance for the following quantities, because these

assert np.allclose(
    [alpha_x],
    [
        0.0,
    ],
    atol=0.1,
)
assert np.allclose(
    [dispersion_x],
    [
        0.0,
    ],
    atol=3.0e-5,
)

Dogleg in Reverse Using Envelope Solver

This is identical to reverse dogleg example, except that envelope tracking is used instead of particle tracking.

The primary purpose is to demonstrate the initialization of a beam with a nonzero x-pt correlation in a lattice region with nonzero dispersion.

In this example, the x-pt correlation of the beam is removed at the dogleg exit, and the dispersion goes to zero.

The initial dispersion is taken to be -267 mm.

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 initial and final values of \(\alpha_x\), \(\alpha_y\), \(\beta_x\), \(\beta_y\), \(D_x\), and \(D_{px}\) must agree with nominal values.

Run

This example can be run either as:

• Python script: python3 run_dogleg_reverse_envelope.py or

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

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

Python: Script

Listing 123 You can copy this file from examples/dogleg/run_dogleg_reverse_envelope.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Marco Garten, Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements, twiss

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 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 5.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

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

#   particle bunch
distr = distribution.Waterbag(
    **twiss(
        beta_x=14.374073092185883,
        beta_y=4.6003616743069093,
        beta_t=1.4153999755977573,
        emitt_x=8.180911194584674375e-11,
        emitt_y=1.020438927769593e-10,
        emitt_t=8.5698645128105327e-09,
        alpha_x=1.3385830279518021,
        alpha_y=-1.3479109197361046,
        alpha_t=92.624347459169869,
        dispersion_x=-0.26669723385388505,
    ),
)
sim.init_envelope(ref, distr, bunch_charge_C)

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

# design the accelerator lattice
ns = 25  # number of slices per ds in the element
rc = 10.3462283686195526  # bend radius (meters)
psi = 0.048345620280243  # pole face rotation angle (radians)
lb = 0.500194828041958  # bend arc length (meters)

# Drift elements
dr1 = elements.Drift(name="dr1", ds=5.0058489435, nslice=ns)
dr2 = elements.Drift(name="dr2", ds=0.5, nslice=ns)

# Bend elements
sbend1 = elements.Sbend(name="sbend1", ds=lb, rc=-rc, nslice=ns)
sbend2 = elements.Sbend(name="sbend2", ds=lb, rc=rc, nslice=ns)

# Dipole Edge Focusing elements
dipedge1 = elements.DipEdge(name="dipedge1", psi=-psi, rc=-rc, g=0.0, K2=0.0)
dipedge2 = elements.DipEdge(name="dipedge2", psi=psi, rc=rc, g=0.0, K2=0.0)

lattice_dogleg = [sbend1, dipedge1, dr1, dipedge2, sbend2, dr2]

sim.lattice.append(monitor)
lattice_dogleg.reverse()
sim.lattice.extend(lattice_dogleg)
sim.lattice.append(monitor)

# run simulation
sim.track_envelope()

# clean shutdown
sim.finalize()

Executable: Input File

Listing 124 You can copy this file from examples/dogleg/input_dogleg_reverse_envelope.in.

###############################################################################
# Particle Beam(s)
###############################################################################
beam.kin_energy = 5.0e3
beam.particle = electron
beam.distribution = waterbag_from_twiss
beam.alphaX = 1.33858302795180
beam.alphaY = -1.347910457923642
beam.alphaT = 92.624347459105241
beam.betaX = 14.37407309218588
beam.betaY = 4.600362059144475
beam.betaT = 1.4153999755967554
beam.emittX = 8.180911194584674375e-11
beam.emittY = 1.020438927769593e-10
beam.emittT = 8.5698645128164239e-09
beam.dispX = -0.2667

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 monitor
lattice.reverse = true
lattice.nslice = 25

sbend1.type = sbend
sbend1.ds = 0.500194828041958       # projected length 0.5 m, angle 2.77 deg
sbend1.rc = -10.3462283686195526

drift1.type = drift
drift1.ds = 5.0058489435            # projected length 5 m

sbend2.type = sbend
sbend2.ds = sbend1.ds               # projected length 0.5 m, angle 2.77 deg
sbend2.rc = -sbend1.rc

drift2.type = drift
drift2.ds = 0.5

dipedge1.type = dipedge   # dipole edge focusing
dipedge1.psi = -0.048345620280243
dipedge1.rc = -10.3462283686195526
dipedge1.g = 0.0
dipedge1.K2 = 0.0

dipedge2.type = dipedge
dipedge2.psi = -dipedge1.psi
dipedge2.rc = -dipedge1.rc
dipedge2.g = 0.0
dipedge2.K2 = 0.0

monitor.type = beam_monitor
monitor.backend = h5


###############################################################################
# Algorithms
###############################################################################
algo.track = "envelope"
algo.space_charge = false


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

Analyze

We run the following script to analyze correctness:

Script analysis_dogleg_reverse_envelope.py

Listing 125 You can copy this file from examples/dogleg/analysis_dogleg_reverse_envelope.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"]
alpha_x = rbc["alpha_x"]
beta_x = rbc["beta_x"]
alpha_y = rbc["alpha_y"]
beta_y = rbc["beta_y"]
dispersion_x = rbc["dispersion_x"]
dispersion_px = rbc["dispersion_px"]
dispersion_y = rbc["dispersion_y"]
dispersion_py = rbc["dispersion_py"]

sigma_xi = sigma_x.iloc[0]
sigma_yi = sigma_y.iloc[0]
sigma_ti = sigma_t.iloc[0]
emittance_xi = emittance_x.iloc[0]
emittance_yi = emittance_y.iloc[0]
emittance_ti = emittance_t.iloc[0]
alpha_xi = alpha_x.iloc[0]
beta_xi = beta_x.iloc[0]
alpha_yi = alpha_y.iloc[0]
beta_yi = beta_y.iloc[0]
dispersion_xi = dispersion_x.iloc[0]
dispersion_pxi = dispersion_px.iloc[0]
dispersion_yi = dispersion_y.iloc[0]
dispersion_pyi = dispersion_py.iloc[0]

length = len(s) - 1

sf = s.iloc[length]
sigma_xf = sigma_x.iloc[length]
sigma_yf = sigma_y.iloc[length]
sigma_tf = sigma_t.iloc[length]
emittance_xf = emittance_x.iloc[length]
emittance_yf = emittance_y.iloc[length]
emittance_tf = emittance_t.iloc[length]
alpha_xf = alpha_x.iloc[length]
beta_xf = beta_x.iloc[length]
alpha_yf = alpha_y.iloc[length]
beta_yf = beta_y.iloc[length]
dispersion_xf = dispersion_x.iloc[length]
dispersion_pxf = dispersion_px.iloc[length]
dispersion_yf = dispersion_y.iloc[length]
dispersion_pyf = dispersion_py.iloc[length]


print("Initial Beam:")
print(f"  sigx={sigma_xi:e} sigy={sigma_yi:e} sigt={sigma_ti:e}")
print(
    f"  emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}"
)

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

assert np.allclose(
    [sigma_xi, sigma_yi, sigma_ti, emittance_xi, emittance_yi, emittance_ti],
    [
        1.924711e-03,
        2.165646e-05,
        1.102534e-04,
        7.668809e-09,
        1.018986e-10,
        8.588054e-09,
    ],
    rtol=rtol,
    atol=atol,
)


print("")
print("Final Beam:")
print(f"  sigx={sigma_xf:e} sigy={sigma_yf:e} sigt={sigma_tf:e}")
print(
    f"  emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}"
)

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

assert np.allclose(
    [sigma_xf, sigma_yf, sigma_tf, emittance_xf, emittance_yf, emittance_tf],
    [
        2.057123e-05,
        6.911405e-05,
        2.012573e-05,
        8.174618e-11,
        1.018986e-10,
        1.151058e-08,
    ],
    rtol=rtol,
    atol=atol,
)


print("")
print("Initial Twiss functions:")
print(
    f"  alpha_x={alpha_xi:e} beta_x={beta_xi:e} alpha_y={alpha_yi:e} beta_y={beta_yi:e}"
)
print(f"  dispersion_x={dispersion_xi:e} dispersion_px={dispersion_pxi:e}")
print(f"  dispersion_y={dispersion_yi:e} dispersion_py={dispersion_pyi:e}")

atol = 0.0  # ignored
rtol = 2.5e-2
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [alpha_xi, beta_xi, alpha_yi, beta_yi, dispersion_xi],
    [
        1.340770e00,
        1.440253e01,
        -1.347747e00,
        4.602637e00,
        -2.667038e-01,
    ],
    rtol=rtol,
    atol=atol,
)
assert np.allclose(
    [dispersion_pxi, dispersion_yi, dispersion_pyi],
    [
        0.0,
        0.0,
        0.0,
    ],
    atol=4.0e-5,
)

print("")
print("Final Twiss functions:")
print(
    f"  alpha_x={alpha_xf:e} beta_x={beta_xf:e} alpha_y={alpha_yf:e} beta_y={beta_yf:e}"
)
print(f"  dispersion_x={dispersion_xf:e} dispersion_px={dispersion_pxf:e}")
print(f"  dispersion_y={dispersion_yf:e} dispersion_py={dispersion_pyf:e}")

atol = 0.0  # ignored
rtol = 2.0e-2
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [beta_xf, alpha_yf, beta_yf],
    [
        5.176642e00,
        -4.973010e00,
        4.687750e01,
    ],
    rtol=rtol,
    atol=atol,
)

# We use absolute tolerance for the following quantities, because these

assert np.allclose(
    [alpha_xf],
    [
        0.0,
    ],
    atol=0.1,
)
assert np.allclose(
    [dispersion_xf],
    [
        0.0,
    ],
    atol=4.0e-5,
)

Vertical Dogleg

This is identical to the dogleg example, except that every bend (and its dip-edge focusing) is rotated by 90 degrees, so the lattice bends in the vertical (\(y\)) plane. The dogleg therefore generates dispersion in \(y\) instead of \(x\).

To obtain the exact \(x \leftrightarrow y\) mirror image of the horizontal dogleg, the initial Twiss functions are exchanged \(x \leftrightarrow y\) as well.

The same lattice is tracked first with the envelope (covariance) solver and then with particle tracking, in the same script.

The script asserts in-situ that the final values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), \(\epsilon_t\), \(\alpha_x\), \(\alpha_y\), \(\beta_x\), \(\beta_y\), and \(D_y\) agree with nominal values – for both the envelope and the particle beam (the dispersion now appears in \(y\)).

Run

This example runs both envelope and particle tracking and is provided as a Python script: python3 run_dogleg_vertical.py.

Listing 126 You can copy this file from examples/dogleg/run_dogleg_vertical.py.

#!/usr/bin/env python3
#
# Copyright 2022-2026 ImpactX contributors
# Authors: Marco Garten, Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

import numpy as np

from impactx import ImpactX, distribution, elements

# This is the vertical-bending variant of the 2-bend dogleg lattice: every bend
# (and its dip-edge focusing) is rolled by 90 degrees, so the dispersion that
# the dogleg generates appears in the vertical (y) plane instead of horizontal
# (x). To obtain the x <-> y mirror image of the horizontal dogleg, the initial
# Twiss functions are exchanged x <-> y as well.
#
# The same lattice is run first with the envelope (covariance) solver and then
# with particle tracking, exercising both the linear-map and the particle path
# through rolled (tilted) elements.

kin_energy_MeV = 5.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles


def dogleg_vertical_distribution():
    """A 5 GeV electron Waterbag with x <-> y exchanged relative to the
    horizontal dogleg (see ``run_dogleg.py``)."""
    return distribution.Waterbag(
        lambdaX=1.3084093142e-5,  # exchanged with lambdaY
        lambdaY=2.2951017632e-5,  # exchanged with lambdaX
        lambdaT=5.5555553e-8,
        lambdaPx=2.803697378e-6,  # exchanged with lambdaPy
        lambdaPy=1.598353425e-6,  # exchanged with lambdaPx
        lambdaPt=2.000000000e-6,
        muxpx=0.933345606203060,  # == muypy, so the exchange leaves it unchanged
        muypy=0.933345606203060,
        mutpt=0.999999961419755,
    )


def dogleg_vertical_lattice():
    """The horizontal dogleg lattice with all bends and dip edges rolled by
    90 degrees, so it bends in the vertical (y) plane."""
    ns = 25  # number of slices per ds in the element
    rc = 10.3462283686195526  # bend radius (meters)
    psi = 0.048345620280243  # pole face rotation angle (radians)
    lb = 0.500194828041958  # bend arc length (meters)
    roll = 90.0  # roll the bending plane from horizontal (x) to vertical (y)

    # Drift elements
    dr1 = elements.Drift(name="dr1", ds=5.0058489435, nslice=ns)
    dr2 = elements.Drift(name="dr2", ds=0.5, nslice=ns)

    # Bend elements (rolled by 90 degrees)
    sbend1 = elements.Sbend(name="sbend1", ds=lb, rc=-rc, rotation=roll, nslice=ns)
    sbend2 = elements.Sbend(name="sbend2", ds=lb, rc=rc, rotation=roll, nslice=ns)

    # Dipole Edge Focusing elements (rolled by 90 degrees)
    dipedge1 = elements.DipEdge(
        name="dipedge1", psi=-psi, rc=-rc, g=0.0, K2=0.0, rotation=roll
    )
    dipedge2 = elements.DipEdge(
        name="dipedge2", psi=psi, rc=rc, g=0.0, K2=0.0, rotation=roll
    )

    return [sbend1, dipedge1, dr1, dipedge2, sbend2, dr2]


def assert_final_beam(moments):
    """Check the final beam against the dogleg reference values: the same
    final-beam properties that the horizontal dogleg checks in
    ``analysis_dogleg.py``, with x <-> y exchanged. The 90-degree roll makes the
    dogleg generate its dispersion in the vertical (y) plane, so it must show up
    in ``dispersion_y`` (and x <-> y throughout). Accepts the moments dict from
    either the envelope or the particle beam."""
    assert np.allclose(
        [
            moments["sigma_x"],
            moments["sigma_y"],
            moments["sigma_t"],
            moments["emittance_x"],
            moments["emittance_y"],
            moments["emittance_t"],
            moments["alpha_x"],
            moments["beta_x"],
            moments["alpha_y"],
            moments["beta_y"],
            moments["dispersion_y"],
        ],
        [
            2.166654e-05,
            1.922660e-03,
            1.101353e-04,
            1.020439e-10,
            8.561046e-09,
            8.569865e-09,
            -1.347910e00,
            4.600362e00,
            1.338583e00,
            1.437407e01,
            -2.666972e-01,
        ],
        rtol=2.0e-2,
    )


# set up the simulation: reference particle, lattice, and diagnostics
sim = ImpactX()
sim.space_charge = False
sim.slice_step_diagnostics = True
sim.init_grids()

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

monitor = elements.BeamMonitor("monitor", backend="h5")
sim.lattice.append(monitor)
sim.lattice.extend(dogleg_vertical_lattice())
sim.lattice.append(monitor)

# The same lattice is run with both solvers; init_envelope() copies the
# reference particle, so envelope tracking does not disturb the particle run.

# (1) envelope (covariance) tracking: exercises the linear-map (transport-map)
# path through the rolled elements.
sim.init_envelope(ref, dogleg_vertical_distribution(), bunch_charge_C)
sim.track_envelope()
assert_final_beam(sim.envelope.beam_moments(ref))

# (2) particle tracking: exercises the particle push through the rolled elements.
sim.add_particles(bunch_charge_C, dogleg_vertical_distribution(), npart)
sim.track_particles()
assert_final_beam(sim.beam.beam_moments())

sim.finalize()

Dogleg with Energy Jitter

This is identical to the dogleg example, except the initial beam distribution has been given a 2.5% offset in the value of mean energy.

The primary purpose is to demonstrate the use of a beam centroid offset to study the effects of, e.g. energy jitter.

The 2.5% energy offset couples through the lattice R16 (dispersion) to result in a mean horizontal x-offset at the end of the dogleg.

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 initial and final values of \(\alpha_x\), \(\alpha_y\), \(\beta_x\), \(\beta_y\), \(D_x\), and \(D_{px}\) must agree with nominal values.

Finally, the values of \(\mean_pt\), \(\mean_x\), and \(\mean_{px}\) must agree with predicted values.

Run

This example can be run either as:

• Python script: python3 run_dogleg_jitter.py or

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

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

Python: Script

Listing 127 You can copy this file from examples/dogleg/run_dogleg_jitter.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Marco Garten, 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 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 5.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

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

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=2.2951017632e-5,
    lambdaY=1.3084093142e-5,
    lambdaT=5.5555553e-8,
    lambdaPx=1.598353425e-6,
    lambdaPy=2.803697378e-6,
    lambdaPt=2.000000000e-6,
    muxpx=0.933345606203060,
    muypy=0.933345606203060,
    mutpt=0.999999961419755,
    meanPt=0.025,
)
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
rc = 10.3462283686195526  # bend radius (meters)
psi = 0.048345620280243  # pole face rotation angle (radians)
lb = 0.500194828041958  # bend arc length (meters)

# Drift elements
dr1 = elements.Drift(name="dr1", ds=5.0058489435, nslice=ns)
dr2 = elements.Drift(name="dr2", ds=0.5, nslice=ns)

# Bend elements
sbend1 = elements.Sbend(name="sbend1", ds=lb, rc=-rc, nslice=ns)
sbend2 = elements.Sbend(name="sbend2", ds=lb, rc=rc, nslice=ns)

# Dipole Edge Focusing elements
dipedge1 = elements.DipEdge(name="dipedge1", psi=-psi, rc=-rc, g=0.0, K2=0.0)
dipedge2 = elements.DipEdge(name="dipedge2", psi=psi, rc=rc, g=0.0, K2=0.0)

lattice_dogleg = [sbend1, dipedge1, dr1, dipedge2, sbend2, dr2]

sim.lattice.append(monitor)
sim.lattice.extend(lattice_dogleg)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Executable: Input File

Listing 128 You can copy this file from examples/dogleg/input_dogleg_jitter.in.

###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 5.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 2.2951017632e-5
beam.lambdaY = 1.3084093142e-5
beam.lambdaT = 5.5555553e-8
beam.lambdaPx = 1.598353425e-6
beam.lambdaPy = 2.803697378e-6
beam.lambdaPt = 2.000000000e-6
beam.muxpx = 0.933345606203060
beam.muypy = beam.muxpx
beam.mutpt = 0.999999961419755
beam.meanPt = 0.025

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 monitor
lattice.nslice = 25

sbend1.type = sbend
sbend1.ds = 0.500194828041958       # projected length 0.5 m, angle 2.77 deg
sbend1.rc = -10.3462283686195526

drift1.type = drift
drift1.ds = 5.0058489435            # projected length 5 m

sbend2.type = sbend
sbend2.ds = sbend1.ds               # projected length 0.5 m, angle 2.77 deg
sbend2.rc = -sbend1.rc

drift2.type = drift
drift2.ds = 0.5

dipedge1.type = dipedge   # dipole edge focusing
dipedge1.psi = -0.048345620280243
dipedge1.rc = -10.3462283686195526
dipedge1.g = 0.0
dipedge1.K2 = 0.0

dipedge2.type = dipedge
dipedge2.psi = -dipedge1.psi
dipedge2.rc = -dipedge1.rc
dipedge2.g = 0.0
dipedge2.K2 = 0.0

monitor.type = beam_monitor
monitor.backend = h5


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


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

Analyze

We run the following script to analyze correctness:

Script analysis_dogleg_jitter.py

Listing 129 You can copy this file from examples/dogleg/analysis_dogleg_jitter.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)


def get_twiss(openpmd_beam):
    """Return Twiss functions from an openPMD particle species

    Returns
    -------
    alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px
    """
    alpha_x = openpmd_beam.get_attribute("alpha_x")
    beta_x = openpmd_beam.get_attribute("beta_x")
    d_x = openpmd_beam.get_attribute("dispersion_x")
    d_px = openpmd_beam.get_attribute("dispersion_px")
    alpha_y = openpmd_beam.get_attribute("alpha_y")
    beta_y = openpmd_beam.get_attribute("beta_y")
    # d_y = openpmd_beam.get_attribute("dispersion_y")
    # d_py = openpmd_beam.get_attribute(["dispersion_py")

    return (alpha_x, beta_x, alpha_y, beta_y, d_x, d_px)


def get_mean(openpmd_beam):
    """Return centroid (mean) from an openPMD particle species

    Returns
    -------
    x_mean, px_mean, y_mean, py_mean, t_mean, pt_mean
    """
    x_mean = openpmd_beam.get_attribute("mean_x")
    px_mean = openpmd_beam.get_attribute("mean_px")
    y_mean = openpmd_beam.get_attribute("mean_y")
    py_mean = openpmd_beam.get_attribute("mean_py")
    t_mean = openpmd_beam.get_attribute("mean_t")
    pt_mean = openpmd_beam.get_attribute("mean_pt")

    return (x_mean, px_mean, y_mean, py_mean, t_mean, pt_mean)


# 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_beam = series.iterations[last_step].particles["beam"]
final = final_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 = 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],
    [
        6.4214719960819659e-005,
        3.6603372435649773e-005,
        1.9955175623579313e-004,
        1.0198263116327677e-010,
        1.0308359092878036e-010,
        4.0035161705244885e-010,
    ],
    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 = 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],
    [
        1.922660e-03,
        2.166654e-05,
        1.101353e-04,
        8.561046e-09,
        1.020439e-10,
        8.569865e-09,
    ],
    rtol=rtol,
    atol=atol,
)

print("")
print("Final Twiss functions:")
alpha_x, beta_x, alpha_y, beta_y, dispersion_x, dispersion_px = get_twiss(final_beam)
print(f"  alpha_x={alpha_x:e} beta_x={beta_x:e} alpha_y={alpha_y:e} beta_y={beta_y:e}")
print(f"  dispersion_x={dispersion_x:e} dispersion_px={dispersion_px: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(
    [alpha_x, beta_x, alpha_y, beta_y, dispersion_x],
    [
        1.338583e00,
        1.437407e01,
        -1.347910e00,
        4.600362e00,
        -2.666972e-01,
    ],
    rtol=rtol,
    atol=atol,
)

print("")
print("Final beam centroid:")
x_mean, px_mean, y_mean, py_mean, t_mean, pt_mean = get_mean(final_beam)
print(f"  x_mean={x_mean:e} y_mean={y_mean:e} t_mean={t_mean:e}")
print(f"  px_mean={px_mean:e} py_mean={py_mean:e} pt_mean={pt_mean:e}")

# Predicted beam centroid
mean_pt_input = 0.025  # specified in the lattice input file
pt_mean_predicted = mean_pt_input
x_mean_predicted = -dispersion_x * pt_mean
px_mean_predicted = -dispersion_px * pt_mean

print("")
print("Predicted beam centroid:")
print(
    f"  x_mean_predicted={x_mean_predicted:e} px_mean_predicted={px_mean_predicted:e} pt_mean_predicted={pt_mean_predicted: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(
    [x_mean, pt_mean],
    [
        x_mean_predicted,
        pt_mean_predicted,
    ],
    rtol=rtol,
    atol=atol,
)
atol = 1.0e-6  # ignored
print(f"  atol={atol}")
assert np.allclose(
    [px_mean],
    [
        px_mean_predicted,
    ],
    atol=atol,
)