Chicane
This is the Berlin-Zeuthen magnetic bunch compression chicane, which is a standardized community benchmark.
All parameters can be found online. A 5 GeV electron bunch with normalized transverse rms emittance of 1 um undergoes longitudinal compression by a factor of 10 in a standard 4-bend chicane.
The emittances should be preserved, and the rms pulse length should decrease by the compression factor (10).
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.
We also have a variation of this test that includes CSR effects in the bending magnets.
Run
This example can be run either as:
Python script:
python3 run_chicane.pyorImpactX executable using an input file:
impactx input_chicane.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.
#!/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=1.0, nslice=ns)
dr3 = elements.Drift(name="dr3", ds=2.0, 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_half = [sbend1, dipedge1, monitor, dr1, monitor, dipedge2, sbend2, monitor]
# assign a segment with the first half of the lattice
sim.lattice.append(monitor)
sim.lattice.extend(lattice_half)
sim.lattice.append(dr2)
lattice_half.reverse()
# extend the lattice by a reversed half
sim.lattice.extend(lattice_half)
sim.lattice.append(monitor)
sim.lattice.append(dr3)
sim.lattice.append(monitor)
# run simulation
sim.track_particles()
# clean shutdown
sim.finalize()
###############################################################################
# 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 monitor drift1 monitor dipedge2 sbend2 monitor drift2 \
monitor sbend2 dipedge2 monitor drift1 monitor dipedge1 sbend1 monitor drift3 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 = 0.500194828041958 # projected length 0.5 m, angle 2.77 deg
sbend2.rc = 10.3462283686195526
drift2.type = drift
drift2.ds = 1.0
drift3.type = drift
drift3.ds = 2.0
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 = 0.048345620280243
dipedge2.rc = 10.3462283686195526
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_chicane.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 # 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],
[
2.3928429374387210e-005,
8.4424535301423173e-005,
1.9976426324802290e-005,
1.0198281373761584e-010,
1.0308356090529235e-010,
4.0027996099961315e-010,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the beam evolution over time:
Script plot_chicane.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
import openpmd_api as io
import pandas as pd
from matplotlib.ticker import MaxNLocator
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
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()
ref_particle = read_time_series("diags/ref_particle.*")
# steps & corresponding z
steps = list(series.iterations)
z = list(
map(lambda step: ref_particle[ref_particle["step"] == step].z.values[0], steps)
)
# 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
# read reduced diagnostics
rbc = read_time_series("diags/reduced_beam_characteristics.*")
s = rbc["s"]
sigma_x = rbc["sigma_x"] * millimeter
sigma_y = rbc["sigma_y"] * millimeter
sigma_t = rbc["sigma_t"] * millimeter
emittance_x = rbc["emittance_x"] * nm_rad
emittance_y = rbc["emittance_y"] * nm_rad
emittance_t = rbc["emittance_t"] * nm_rad
length = len(s) - 1
# select a single particle by id
# particle_42 = beam[beam["id"] == 42]
# print(particle_42)
# steps & corresponding z
steps = list(series.iterations)
# print beam transverse size over steps
f = plt.figure(figsize=(9, 4.8))
ax1 = f.gca()
im_sigx = ax1.plot(s, sigma_x, label=r"$\sigma_x$")
im_sigt = ax1.plot(s, sigma_t, label=r"$\sigma_t$")
ax2 = ax1.twinx()
ax2.set_prop_cycle(None) # reset color cycle
im_emittance_x = ax2.plot(s, emittance_x, ":", label=r"$\epsilon_x$")
im_emittance_t = ax2.plot(s, emittance_t, ":", label=r"$\epsilon_t$")
ax1.legend(
handles=im_sigx + im_sigt + im_emittance_x + im_emittance_t, loc="upper right"
)
ax1.set_xlabel(r"$z$ [m]")
ax1.set_ylabel(r"$\sigma_{x,t}$ [mm]")
ax2.set_ylabel(r"$\epsilon_{x,t}$ [nm]")
ax1.set_ylim([0.0, 2.0])
ax2.set_ylim([0.0, 23.0])
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
num_plots_per_row = len(steps)
fig, axs = plt.subplots(
4, num_plots_per_row, figsize=(9, 6.4), sharex="row", sharey="row"
)
ncol_ax = -1
for step in steps:
# plot initial distribution & at exit of each element
ncol_ax += 1
# t-pt
ax = axs[(0, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_t.multiply(millimeter),
beam_at_step.momentum_t.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$ct$ [mm]")
z_unit = ""
if ncol_ax == num_plots_per_row - 1:
z_unit = " [m]"
ax.set_title(f"$z={z[ncol_ax]:.1f}${z_unit}")
# x-px
ax = axs[(1, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_x.multiply(millimeter),
beam_at_step.momentum_x.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$x$ [mm]")
# t-x
ax = axs[(2, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_t.multiply(millimeter),
beam_at_step.position_x.multiply(millimeter),
s=0.01,
)
ax.set_xlabel(r"$ct$ [mm]")
# t-px
ax = axs[(3, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_t.multiply(millimeter),
beam_at_step.momentum_x.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$ct$ [mm]")
axs[(0, 0)].set_ylabel(r"$p_t$ [mrad]")
axs[(1, 0)].set_ylabel(r"$p_x$ [mrad]")
axs[(2, 0)].set_ylabel(r"$x$ [mm]")
axs[(3, 0)].set_ylabel(r"$p_x$ [mrad]")
plt.tight_layout()
if args.save_png:
plt.savefig("chicane_scatter.png")
else:
plt.show()
Fig. 3 (top) Chicane floorplan. (bottom) Chicane beam width and emittance evolution.
Fig. 4 Chicane beam width and emittance evolution