The “bare” linear lattice of the Fermilab IOTA storage ring

The linear lattice of the IOTA storage ring, configured for operation with a 2.5 MeV proton beam.

The drift regions available for insertion of the special nonlinear magnetic element for integrable optics experiments are denoted dnll.

The second moments of the particle distribution after a single turn should coincide with the initial section moments of the particle distribution, to within the level expected due to numerical particle noise. The example runs 5 turns.

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

Run

This example can be run either as:

  • Python script: python3 run_iotalattice.py or

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

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

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

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

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = False
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = True

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

# init particle beam
kin_energy_MeV = 2.5
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000

#   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)

#   particle bunch
distr = distribution.Waterbag(
    sigmaX=1.588960728035e-3,
    sigmaY=2.496625268437e-3,
    sigmaT=1.0e-3,
    sigmaPx=2.8320397837724e-3,
    sigmaPy=1.802433091137e-3,
    sigmaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

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

# init accelerator lattice
ns = 10  # number of slices per ds in the element

# Drift elements
dra1 = elements.Drift(ds=0.9125, nslice=ns)
dra2 = elements.Drift(ds=0.135, nslice=ns)
dra3 = elements.Drift(ds=0.725, nslice=ns)
dra4 = elements.Drift(ds=0.145, nslice=ns)
dra5 = elements.Drift(ds=0.3405, nslice=ns)
drb1 = elements.Drift(ds=0.3205, nslice=ns)
drb2 = elements.Drift(ds=0.14, nslice=ns)
drb3 = elements.Drift(ds=0.1525, nslice=ns)
drb4 = elements.Drift(ds=0.31437095, nslice=ns)
drc1 = elements.Drift(ds=0.42437095, nslice=ns)
drc2 = elements.Drift(ds=0.355, nslice=ns)
dnll = elements.Drift(ds=1.8, nslice=ns)
drd1 = elements.Drift(ds=0.62437095, nslice=ns)
drd2 = elements.Drift(ds=0.42, nslice=ns)
drd3 = elements.Drift(ds=1.625, nslice=ns)
drd4 = elements.Drift(ds=0.6305, nslice=ns)
dre1 = elements.Drift(ds=0.5305, nslice=ns)
dre2 = elements.Drift(ds=1.235, nslice=ns)
dre3 = elements.Drift(ds=0.8075, nslice=ns)

# Bend elements
rc30 = 0.822230996255981
sbend30 = elements.Sbend(ds=0.4305191429, rc=rc30)
edge30 = elements.DipEdge(psi=0.0, rc=rc30, g=0.058, K2=0.5)

rc60 = 0.772821121503940
sbend60 = elements.Sbend(ds=0.8092963858, rc=rc60)
edge60 = elements.DipEdge(psi=0.0, rc=rc60, g=0.058, K2=0.5)

# Quad elements
ds_quad = 0.21
qa1 = elements.Quad(ds=ds_quad, k=-8.78017699, nslice=ns)
qa2 = elements.Quad(ds=ds_quad, k=13.24451745, nslice=ns)
qa3 = elements.Quad(ds=ds_quad, k=-13.65151327, nslice=ns)
qa4 = elements.Quad(ds=ds_quad, k=19.75138652, nslice=ns)
qb1 = elements.Quad(ds=ds_quad, k=-10.84199727, nslice=ns)
qb2 = elements.Quad(ds=ds_quad, k=16.24844348, nslice=ns)
qb3 = elements.Quad(ds=ds_quad, k=-8.27411104, nslice=ns)
qb4 = elements.Quad(ds=ds_quad, k=-7.45719247, nslice=ns)
qb5 = elements.Quad(ds=ds_quad, k=14.03362243, nslice=ns)
qb6 = elements.Quad(ds=ds_quad, k=-12.23595641, nslice=ns)
qc1 = elements.Quad(ds=ds_quad, k=-13.18863768, nslice=ns)
qc2 = elements.Quad(ds=ds_quad, k=11.50601829, nslice=ns)
qc3 = elements.Quad(ds=ds_quad, k=-11.10445869, nslice=ns)
qd1 = elements.Quad(ds=ds_quad, k=-6.78179218, nslice=ns)
qd2 = elements.Quad(ds=ds_quad, k=5.19026998, nslice=ns)
qd3 = elements.Quad(ds=ds_quad, k=-5.8586173, nslice=ns)
qd4 = elements.Quad(ds=ds_quad, k=4.62460039, nslice=ns)
qe1 = elements.Quad(ds=ds_quad, k=-4.49607687, nslice=ns)
qe2 = elements.Quad(ds=ds_quad, k=6.66737146, nslice=ns)
qe3 = elements.Quad(ds=ds_quad, k=-6.69148177, nslice=ns)

# build lattice: first half, qe3, then mirror
# fmt: off
lattice_half = [
    dra1, qa1, dra2, qa2, dra3, qa3, dra4, qa4, dra5,
    edge30, sbend30, edge30, drb1, qb1, drb2, qb2, drb2, qb3,
    drb3, dnll, drb3, qb4, drb2, qb5, drb2, qb6, drb4,
    edge60, sbend60, edge60, drc1, qc1, drc2, qc2, drc2, qc3, drc1,
    edge60, sbend60, edge60, drd1, qd1, drd2, qd2, drd3, qd3, drd2, qd4, drd4,
    edge30, sbend30, edge30, dre1, qe1, dre2, qe2, dre3
]
# fmt:on
sim.lattice.append(monitor)
sim.lattice.extend(lattice_half)
sim.lattice.append(qe3)
lattice_half.reverse()
sim.lattice.extend(lattice_half)
sim.lattice.append(monitor)

# number of turns in the ring
sim.periods = 5

# run simulation
sim.evolve()

# clean shutdown
del sim
amr.finalize()
Listing 40 You can copy this file from examples/iota_lattice/input_iotalattice.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.sigmaX = 1.588960728035e-3
beam.sigmaY = 2.496625268437e-3
beam.sigmaT = 1.0e-3
beam.sigmaPx = 2.8320397837724e-3
beam.sigmaPy = 1.802433091137e-3
beam.sigmaPt = 0.0
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.periods = 5
lattice.elements = monitor first_half qe3 second_half monitor

# lines
first_half.type = line
first_half.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5                \
                      edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3        \
                      drb3 dnll drb3 qb4 drb2 qb5 drb2 qb6 drb4               \
                      edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1   \
                      edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4  \
                      edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3

second_half.type = line
second_half.reverse = true
second_half.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5               \
                       edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3       \
                       drb3 dnll drb3 qb4 drb2 qb5 drb2 qb6 drb4              \
                       edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1  \
                       edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4  \
                       edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3

# thick element splitting for space charge
lattice.nslice = 10


# Drift elements:

dra1.type = drift
dra1.ds = 0.9125

dra2.type = drift
dra2.ds = 0.135

dra3.type = drift
dra3.ds = 0.725

dra4.type = drift
dra4.ds = 0.145

dra5.type = drift
dra5.ds = 0.3405

drb1.type = drift
drb1.ds = 0.3205

drb2.type = drift
drb2.ds = 0.14

drb3.type = drift
drb3.ds = 0.1525

drb4.type = drift
drb4.ds = 0.31437095

drc1.type = drift
drc1.ds = 0.42437095

drc2.type = drift
drc2.ds = 0.355

dnll.type = drift
dnll.ds = 1.8

drd1.type = drift
drd1.ds = 0.62437095

drd2.type = drift
drd2.ds = 0.42

drd3.type = drift
drd3.ds = 1.625

drd4.type = drift
drd4.ds = 0.6305

dre1.type = drift
dre1.ds = 0.5305

dre2.type = drift
dre2.ds = 1.235

dre3.type = drift
dre3.ds = 0.8075


# Bend elements:

sbend30.type = sbend
sbend30.ds = 0.4305191429
sbend30.rc = 0.822230996255981

edge30.type = dipedge
edge30.psi = 0.0
edge30.rc = 0.822230996255981
edge30.g = 0.058
edge30.K2 = 0.5

sbend60.type = sbend
sbend60.ds = 0.8092963858
sbend60.rc = 0.772821121503940

edge60.type = dipedge
edge60.psi = 0.0
edge60.rc = 0.772821121503940
edge60.g = 0.058
edge60.K2 = 0.5


# Quad elements:

qa1.type = quad
qa1.ds = 0.21
qa1.k = -8.78017699

qa2.type = quad
qa2.ds = 0.21
qa2.k = 13.24451745

qa3.type = quad
qa3.ds = 0.21
qa3.k = -13.65151327

qa4.type = quad
qa4.ds = 0.21
qa4.k = 19.75138652

qb1.type = quad
qb1.ds = 0.21
qb1.k = -10.84199727

qb2.type = quad
qb2.ds = 0.21
qb2.k = 16.24844348

qb3.type = quad
qb3.ds = 0.21
qb3.k = -8.27411104

qb4.type = quad
qb4.ds = 0.21
qb4.k = -7.45719247

qb5.type = quad
qb5.ds = 0.21
qb5.k = 14.03362243

qb6.type = quad
qb6.ds = 0.21
qb6.k = -12.23595641

qc1.type = quad
qc1.ds = 0.21
qc1.k = -13.18863768

qc2.type = quad
qc2.ds = 0.21
qc2.k = 11.50601829

qc3.type = quad
qc3.ds = 0.21
qc3.k = -11.10445869

qd1.type = quad
qd1.ds = 0.21
qd1.k = -6.78179218

qd2.type = quad
qd2.ds = 0.21
qd2.k = 5.19026998

qd3.type = quad
qd3.ds = 0.21
qd3.k = -5.8586173

qd4.type = quad
qd4.ds = 0.21
qd4.k = 4.62460039

qe1.type = quad
qe1.ds = 0.21
qe1.k = -4.49607687

qe2.type = quad
qe2.ds = 0.21
qe2.k = 6.66737146

qe3.type = quad
qe3.ds = 0.21
qe3.k = -6.69148177

# Beam Monitor: Diagnostics
monitor.type = beam_monitor
monitor.backend = h5


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


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

Analyze

We run the following script to analyze correctness:

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

import numpy as np
import openpmd_api as io
from scipy.stats import moment


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

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

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

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


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

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

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

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

assert np.allclose(
    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
    [1.595934e-03, 2.507263e-03, 9.977588e-04, 4.490896e-06, 4.539378e-06, 0.000000e00],
    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  # a big number
rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, emittance_x, emittance_y, emittance_t],
    [
        1.579848e-03,
        2.510900e-03,
        4.490897e-06,
        4.539378e-06,
        0.0,
    ],
    rtol=rtol,
    atol=atol,
)