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.

Python: Script

Listing 100 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 -*-

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = False
# sim.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(
    lambdaX=1.588960728035e-3,
    lambdaY=2.496625268437e-3,
    lambdaT=1.0e-3,
    lambdaPx=2.8320397837724e-3,
    lambdaPy=1.802433091137e-3,
    lambdaPt=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
sim.finalize()

Executable: Input File

Listing 101 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.lambdaX = 1.588960728035e-3
beam.lambdaY = 2.496625268437e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.8320397837724e-3
beam.lambdaPy = 1.802433091137e-3
beam.lambdaPt = 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
monitor.nonlinear_lens_invariants = true


###############################################################################
# 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 102 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.9 * 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.9 * 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,
)

The full nonlinear lattice of the Fermilab IOTA storage ring

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

The special nonlinear magnetic element for integrable optics experiments is denoted nll. To simplify analysis, the lattice has been arranged so that nll appears as the first element in the sequence.

The two functions H (Hamiltonian) and I (the second invariant) are evaluated at the entrance to the nonlinear element. These values should be unchanged for all particles (to within acceptable tolerance), over the specified number of periods (default 5).

In this test, the initial and final values of \(\mu_H\), \(\sigma_H\), \(\mu_I\), \(\sigma_I\) must agree with nominal values.

Run

This example can be run either as:

• Python script: python3 run_iotalattice_sdep.py or

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

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

Python: Script

Listing 103 You can copy this file from examples/iota_lattice/run_iotalattice_sdep.py.

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

import math

from impactx import ImpactX, distribution, elements

sim = ImpactX()

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

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

# init particle beam
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(energy_MeV)

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=1.865379469388e-003,
    lambdaY=2.0192133150418e-003,
    lambdaT=1.0e-3,
    lambdaPx=1.402566720991e-003,
    lambdaPy=9.57593913381e-004,
    lambdaPt=0.0,
    muxpx=0.482260919078473,
    muypy=0.762127656873158,
    mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
monitor.nonlinear_lens_invariants = True
monitor.alpha = 1.376381920471173
monitor.beta = 1.892632003628881
monitor.tn = 0.4
monitor.cn = 0.01

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

# Special (elliptic) nonlinear element:

# set basic parameters of the nonlinear element
lens_length = 1.8
num_lenses = 18
tune_advance = 0.3
c_parameter = 0.01
t_strength = 0.4
ds = lens_length / num_lenses

# build up the nonlinear lens in segments
ds_half = ds / 2.0
dr = elements.Drift(ds=ds_half, nslice=1)
nll = []
for j in range(0, num_lenses):
    s = -lens_length / 2.0 + ds_half + j * ds
    beta_star = lens_length / 2.0 * 1.0 / math.tan(math.pi * tune_advance)
    beta = beta_star * (
        1.0 + (2.0 * s * math.tan(math.pi * tune_advance) / lens_length) ** 2
    )
    knll_s = t_strength * c_parameter**2 * ds / beta
    cnll_s = c_parameter * math.sqrt(beta)
    nllens = elements.NonlinearLens(knll=knll_s, cnll=cnll_s)
    segment = [dr, nllens, dr]
    nll.extend(segment)

lattice_before_nll = [
    dra1,
    qa1,
    dra2,
    qa2,
    dra3,
    qa3,
    dra4,
    qa4,
    dra5,
    edge30,
    sbend30,
    edge30,
    drb1,
    qb1,
    drb2,
    qb2,
    drb2,
    qb3,
    drb3,
]

lattice_after_nll = [
    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,
]

# build lattice: first half, qe3, then mirror
# modified to place nll as the first element
# fmt:on
sim.lattice.append(monitor)
sim.lattice.extend(nll)
sim.lattice.extend(lattice_after_nll)
sim.lattice.append(qe3)
lattice_after_nll.reverse()
sim.lattice.extend(lattice_after_nll)
sim.lattice.append(dnll)
lattice_before_nll.reverse()
sim.lattice.extend(lattice_before_nll)
lattice_before_nll.reverse()
sim.lattice.extend(lattice_before_nll)
sim.lattice.append(monitor)

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

# run simulation
sim.evolve()

# clean shutdown
sim.finalize()

Executable: Input File

Listing 104 You can copy this file from examples/iota_lattice/input_iotalattice_sdep.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.lambdaX = 1.865379469388e-003
beam.lambdaY = 2.0192133150418e-003
beam.lambdaT = 1.0e-3
beam.lambdaPx = 1.402566720991e-003
beam.lambdaPy = 9.57593913381e-004
beam.lambdaPt = 0.0
beam.muxpx = 0.482260919078473
beam.muypy = 0.762127656873158
beam.mutpt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.periods = 5
lattice.elements = monitor nll after_nll qe3 after_nll_rev dnll before_nll_rev  \
                   before_nll monitor

# lines
before_nll.type = line
before_nll.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5                \
                      edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3        \
                      drb3

after_nll.type = line
after_nll.elements = 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

before_nll_rev.type = line
before_nll_rev.reverse = true
before_nll_rev.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5                \
                      edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3        \
                      drb3

after_nll_rev.type = line
after_nll_rev.reverse = true
after_nll_rev.elements = 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

nll.type = line
nll.elements = = dr_end nll1 dr nll2 dr nll3 dr nll4 dr nll5 dr nll6  \
                       dr nll7 dr nll8 dr nll9 dr nll9 dr nll8 dr nll7  \
                       dr nll6 dr nll5 dr nll4 dr nll3 dr nll2 dr nll1  \
                       dr_end

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

# Nonlinear lens segments
nll1.type = nonlinear_lens
nll1.knll = 2.2742558121e-6
nll1.cnll = 0.013262040169952

nll2.type = nonlinear_lens
nll2.knll = 2.641786366e-6
nll2.cnll = 0.012304986694423

nll3.type = nonlinear_lens
nll3.knll = 3.076868353e-6
nll3.cnll = 0.011401855643727

nll4.type = nonlinear_lens
nll4.knll = 3.582606522e-6
nll4.cnll = 0.010566482535866

nll5.type = nonlinear_lens
nll5.knll = 4.151211157e-6
nll5.cnll = 0.009816181575902

nll6.type = nonlinear_lens
nll6.knll = 4.754946985e-6
nll6.cnll = 0.0091718544892154

nll7.type = nonlinear_lens
nll7.knll = 5.337102374e-6
nll7.cnll = 0.008657195579489

nll8.type = nonlinear_lens
nll8.knll = 5.811437818e-6
nll8.cnll = 0.008296371635942

nll9.type = nonlinear_lens
nll9.knll = 6.081693334e-6
nll9.cnll = 0.008109941789663


# Drift elements:

dr.type = drift
dr.ds = 0.1

dr_end.type = drift
dr_end.ds = 0.05

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
monitor.nonlinear_lens_invariants = true
monitor.alpha = 1.376381920471173
monitor.beta = 1.892632003628881
monitor.tn = 0.4
monitor.cn = 0.01


###############################################################################
# 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_sdep.py

Listing 105 You can copy this file from examples/iota_lattice/analysis_iotalattice_sdep.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 mean and std dev of functions defining the IOTA invariants
    Returns
    -------
    meanH, sigH, meanI, sigI
    """
    meanH = np.mean(beam["H"])
    sigH = moment(beam["H"], moment=2) ** 0.5
    meanI = np.mean(beam["I"])
    sigI = moment(beam["I"], moment=2) ** 0.5

    return (meanH, sigH, meanI, sigI)


# 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:")
meanH, sigH, meanI, sigI = get_moments(initial)
print(f"  meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")

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

assert np.allclose(
    [meanH, sigH, meanI, sigI],
    [7.263202e-02, 4.454371e-02, 9.288060e-02, 8.211506e-02],
    rtol=rtol,
    atol=atol,
)


print("")
print("Final Beam:")
meanH, sigH, meanI, sigI = get_moments(final)
print(f"  meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")

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

assert np.allclose(
    [meanH, sigH, meanI, sigI],
    [7.263202e-02, 4.454371e-02, 9.288060e-02, 8.211506e-02],
    rtol=rtol,
    atol=atol,
)

# join tables on particle ID, so we can compare the same particle initial->final
beam_joined = final.join(initial, lsuffix="_final", rsuffix="_initial")
# add new columns: dH and dI
beam_joined["dH"] = (beam_joined["H_initial"] - beam_joined["H_final"]).abs()
beam_joined["dI"] = (beam_joined["I_initial"] - beam_joined["I_final"]).abs()
# print(beam_joined)

# particle-wise comparison of H & I initial to final
Hrms = np.sqrt(sigH**2 + meanH**2)
Irms = np.sqrt(sigI**2 + meanI**2)

atol = 4.2e-3 * Hrms
rtol = 0.0  # large number
print()
print(f"  atol={atol} (ignored: rtol~={rtol})")
print(f"  dH_max={beam_joined['dH'].max()}")
assert np.allclose(beam_joined["dH"], 0.0, rtol=rtol, atol=atol)

atol = 5.7e-3 * Irms
rtol = 0.0
print()
print(f"  atol={atol} (ignored: rtol~={rtol})")
print(f"  dI_max={beam_joined['dI'].max()}")
assert np.allclose(beam_joined["dI"], 0.0, rtol=rtol, atol=atol)