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.
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 as a Python script (python3 run_iotalattice.py) or with an app with an input file (impactx input_iotalattice.in).
Each can also be prefixed with an MPI executor, such as mpiexec -n 4 ... or srun -n 4 ..., depending on the system.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import amrex
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
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_energy_MeV(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)
# 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.extend(lattice_half)
sim.lattice.append(qe3)
lattice_half.reverse()
sim.lattice.extend(lattice_half)
# run simulation
sim.evolve()
# clean shutdown
del sim
amrex.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.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.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 qe3
dre3 qe2 dre2 qe1 dre1 edge30 sbend30 edge30
drd4 qd4 drd2 qd3 drd3 qd2 drd2 qd1 drd1 edge60 sbend60 edge60
drc1 qc3 drc2 qc2 drc2 qc1 drc1 edge60 sbend60 edge60
drb4 qb6 drb2 qb5 drb2 qb4 drb3 dnll drb3
qb3 drb2 qb2 drb2 qb1 drb1 edge30 sbend30 edge30
dra5 qa4 dra4 qa3 dra3 qa2 dra2 qa1 dra1
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
###############################################################################
# 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
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import glob
import numpy as np
import pandas as pd
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["x"], moment=2) ** 0.5 # variance -> std dev.
sigpx = moment(beam["px"], moment=2) ** 0.5
sigy = moment(beam["y"], moment=2) ** 0.5
sigpy = moment(beam["py"], moment=2) ** 0.5
sigt = moment(beam["t"], moment=2) ** 0.5
sigpt = moment(beam["pt"], moment=2) ** 0.5
epstrms = beam.cov(ddof=0)
emittance_x = (sigx**2 * sigpx**2 - epstrms["x"]["px"] ** 2) ** 0.5
emittance_y = (sigy**2 * sigpy**2 - epstrms["y"]["py"] ** 2) ** 0.5
emittance_t = (sigt**2 * sigpt**2 - epstrms["t"]["pt"] ** 2) ** 0.5
return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
def read_all_files(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(
(
pd.read_csv(filename, delimiter=r"\s+")
for filename in glob.glob(file_pattern)
),
axis=0,
ignore_index=True,
).set_index("id")
# initial/final beam on rank zero
initial = read_all_files("diags/beam_000000.*")
final = read_all_files("diags/beam_final.*")
# 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.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.579848e-03,
2.510900e-03,
1.208202e-02,
4.490897e-06,
4.539378e-06,
0.0,
],
rtol=rtol,
atol=atol,
)