FODO Cell
Stable FODO cell with a zero-current phase advance of 67.8 degrees.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 2.82161941\) m
\(\alpha_\mathrm{x} = -1.59050035\)
\(\beta_\mathrm{y} = 2.82161941\) m
\(\alpha_\mathrm{y} = 1.59050035\)
We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling.
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_fodo.py) or with an app with an input file (impactx input_fodo.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: Axel Huebl, Chad Mitchell
# 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()
# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
energy_MeV = 2.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_energy_MeV(energy_MeV)
# particle bunch
distr = distribution.Waterbag(
sigmaX=3.9984884770e-5,
sigmaY=3.9984884770e-5,
sigmaT=1.0e-3,
sigmaPx=2.6623538760e-5,
sigmaPy=2.6623538760e-5,
sigmaPt=2.0e-3,
muxpx=-0.846574929020762,
muypy=0.846574929020762,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
ns = 25 # number of slices per ds in the element
fodo = [
elements.Drift(ds=0.25, nslice=ns),
elements.Quad(ds=1.0, k=1.0, nslice=ns),
elements.Drift(ds=0.5, nslice=ns),
elements.Quad(ds=1.0, k=-1.0, nslice=ns),
elements.Drift(ds=0.25, nslice=ns),
]
# assign a fodo segment
sim.lattice.extend(fodo)
# run simulation
sim.evolve()
# clean shutdown
del sim
amrex.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.sigmaX = 3.9984884770e-5
beam.sigmaY = 3.9984884770e-5
beam.sigmaT = 1.0e-3
beam.sigmaPx = 2.6623538760e-5
beam.sigmaPy = 2.6623538760e-5
beam.sigmaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = 0.846574929020762
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = drift1 quad1 drift2 quad2 drift3
lattice.nslice = 25
drift1.type = drift
drift1.ds = 0.25
quad1.type = quad
quad1.ds = 1.0
quad1.k = 1.0
drift2.type = drift
drift2.ds = 0.5
quad2.type = quad
quad2.ds = 1.0
quad2.k = -1.0
drift3.type = drift
drift3.ds = 0.25
###############################################################################
# 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_fodo.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 # ignored
rtol = 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],
[
7.5451170454175073e-005,
7.5441588239210947e-005,
9.9775878164077539e-004,
1.9959540393751392e-009,
2.0175015289132990e-009,
2.0013820193294972e-006,
],
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 = 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],
[
7.4790118496224206e-005,
7.5357525169680140e-005,
9.9775879288128088e-004,
1.9959539836392703e-009,
2.0175014668882125e-009,
2.0013820380883801e-006,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the beam evolution over time:
Script plot_fodo.py
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import argparse
import glob
import re
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
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_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')
# options to run this script
parser = argparse.ArgumentParser(description="Plot the FODO benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()
# initial/final beam on rank zero
beam = read_time_series("diags/beam_[0-9]*.*")
ref_particle = read_time_series("diags/ref_particle.*")
# print(beam)
# print(ref_particle)
# scaling to units
millimeter = 1.0e3 # m->mm
mrad = 1.0e3 # ImpactX uses "static units": momenta are normalized by the magnitude of the momentum of the reference particle p0: px/p0 (rad)
# mm_mrad = 1.e6
nm_rad = 1.0e9
# select a single particle by id
# particle_42 = beam[beam["id"] == 42]
# print(particle_42)
# steps & corresponding z
steps = beam.step.unique()
steps.sort()
# print(f"steps={steps}")
z = list(
map(lambda step: ref_particle[ref_particle["step"] == step].z.values[0], steps)
)
# print(f"z={z}")
# beam transversal size & emittance over steps
moments = list(map(lambda step: (step, get_moments(beam[beam["step"] == step])), steps))
# print(moments)
sigx = list(map(lambda step_val: step_val[1][0] * millimeter, moments))
sigy = list(map(lambda step_val: step_val[1][1] * millimeter, moments))
emittance_x = list(map(lambda step_val: step_val[1][3] * nm_rad, moments))
emittance_y = list(map(lambda step_val: step_val[1][4] * nm_rad, moments))
# print(sigx, sigy)
# print beam transversal size over steps
f = plt.figure(figsize=(9, 4.8))
ax1 = f.gca()
im_sigx = ax1.plot(z, sigx, label=r"$\sigma_x$")
im_sigy = ax1.plot(z, sigy, label=r"$\sigma_y$")
ax2 = ax1.twinx()
ax2._get_lines.prop_cycler = ax1._get_lines.prop_cycler
im_emittance_x = ax2.plot(z, emittance_x, ":", label=r"$\epsilon_x$")
im_emittance_y = ax2.plot(z, emittance_y, ":", label=r"$\epsilon_y$")
ax1.legend(
handles=im_sigx + im_sigy + im_emittance_x + im_emittance_y, loc="lower center"
)
ax1.set_xlabel(r"$z$ [m]")
ax1.set_ylabel(r"$\sigma_{x,y}$ [mm]")
# ax2.set_ylabel(r"$\epsilon_{x,y}$ [mm-mrad]")
ax2.set_ylabel(r"$\epsilon_{x,y}$ [nm]")
ax2.set_ylim([1.5, 2.5])
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
plt.savefig("fodo_sigma.png")
else:
plt.show()
# beam transversal scatter plot over steps
plot_ever_nstep = 25 # this is lattice.nslice in inputs
num_plots_per_row = len(steps) // plot_ever_nstep + 1
fig, axs = plt.subplots(
3, num_plots_per_row, figsize=(9, 4.8), sharex="row", sharey="row"
)
ncol_ax = -1
for step in steps:
# plot initial distribution & at exit of each element
if step != 0 and step % plot_ever_nstep != 0:
continue
print(step)
ncol_ax += 1
# x-y
ax = axs[(0, ncol_ax)]
beam_at_step = beam[beam["step"] == step]
ax.scatter(
beam_at_step.x.multiply(millimeter), beam_at_step.y.multiply(millimeter), s=0.01
)
ax.set_title(f"$z={z[step]}$ [m]")
ax.set_xlabel(r"$x$ [mm]")
# x-px
ax = axs[(1, ncol_ax)]
beam_at_step = beam[beam["step"] == step]
ax.scatter(
beam_at_step.x.multiply(millimeter), beam_at_step.px.multiply(mrad), s=0.01
)
ax.set_xlabel(r"$x$ [mm]")
# y-py
ax = axs[(2, ncol_ax)]
beam_at_step = beam[beam["step"] == step]
ax.scatter(
beam_at_step.y.multiply(millimeter), beam_at_step.py.multiply(mrad), s=0.01
)
ax.set_xlabel(r"$y$ [mm]")
axs[(0, 0)].set_ylabel(r"$y$ [mm]")
axs[(1, 0)].set_ylabel(r"$p_x$ [mrad]")
axs[(2, 0)].set_ylabel(r"$p_y$ [mrad]")
plt.tight_layout()
if args.save_png:
plt.savefig("fodo_scatter.png")
else:
plt.show()
Fig. 1 FODO transversal beam width and emittance evolution
Fig. 2 FODO transversal beam width and phase space evolution