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\), \(\dispersion_x\), and \(\dispersion_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.

Listing 97 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.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=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()
Listing 98 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 99 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\), \(\dispersion_x\), and \(\dispersion_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.

Listing 100 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.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(
    **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()
Listing 101 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 102 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 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\), \(\dispersion_x\), and \(\dispersion_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.

Listing 103 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.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=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()
Listing 104 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 105 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,
)