Chicane
Berlin-Zeuthen magnetic bunch compression chicane: https://www.desy.de/csr/
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.
Run
This example can be run as a Python script (python3 run_chicane.py) or with an app with an input file (impactx input_chicane.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 ImpactX contributors
# Authors: Marco Garten, 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 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
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_energy_MeV(energy_MeV)
# particle bunch
distr = distribution.Waterbag(
sigmaX=2.2951017632e-5,
sigmaY=1.3084093142e-5,
sigmaT=5.5555553e-8,
sigmaPx=1.598353425e-6,
sigmaPy=2.803697378e-6,
sigmaPt=2.000000000e-6,
muxpx=0.933345606203060,
muypy=0.933345606203060,
mutpt=0.999999961419755,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
ns = 25 # number of slices per ds in the element
rc = 10.35 # bend radius (meters)
psi = 0.048345620280243 # pole face rotation angle (radians)
# Drift elements
dr1 = elements.Drift(ds=5.0058489435, nslice=ns)
dr2 = elements.Drift(ds=1.0, nslice=ns)
dr3 = elements.Drift(ds=2.0, nslice=ns)
# Bend elements
sbend1 = elements.Sbend(ds=0.50037, rc=-rc, nslice=ns)
sbend2 = elements.Sbend(ds=0.50037, rc=rc, nslice=ns)
# Dipole Edge Focusing elements
dipedge1 = elements.DipEdge(psi=-psi, rc=-rc, g=0.0, K2=0.0)
dipedge2 = elements.DipEdge(psi=psi, rc=rc, g=0.0, K2=0.0)
lattice_half = [sbend1, dipedge1, dr1, dipedge2, sbend2]
# assign a segment with the first half of the lattice
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(dr3)
# run simulation
sim.evolve()
# clean shutdown
del sim
amrex.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.energy = 5.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.sigmaX = 2.2951017632e-5
beam.sigmaY = 1.3084093142e-5
beam.sigmaT = 5.5555553e-8
beam.sigmaPx = 1.598353425e-6
beam.sigmaPy = 2.803697378e-6
beam.sigmaPt = 2.000000000e-6
beam.muxpx = 0.933345606203060
beam.muypy = 0.933345606203060
beam.mutpt = 0.999999961419755
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 sbend2
dipedge2 drift1 dipedge1 sbend1 drift3
lattice.nslice = 25
sbend1.type = sbend
sbend1.ds = 0.50037 # projected length 0.5 m, angle 2.77 deg
sbend1.rc = -10.35
drift1.type = drift
drift1.ds = 5.0058489435 # projected length 5 m
sbend2.type = sbend
sbend2.ds = 0.50037 # projected length 0.5 m, angle 2.77 deg
sbend2.rc = 10.35
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.35
dipedge1.g = 0.0
dipedge1.K2 = 0.0
dipedge2.type = dipedge
dipedge2.psi = 0.048345620280243
dipedge2.rc = 10.35
dipedge2.g = 0.0
dipedge2.K2 = 0.0
###############################################################################
# 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_chicane.py
#!/usr/bin/env python3
#
# Copyright 2022 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],
[
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 = 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 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 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_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 chicane 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
# for "t": the time coordinate is scaled by c, and therefore has units of length (m) by default, so we can label the axis ct (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)
)
x = list(
map(lambda step: ref_particle[ref_particle["step"] == step].x.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))
sigt = list(map(lambda step_val: step_val[1][2] * millimeter, moments))
emittance_x = list(map(lambda step_val: step_val[1][3] * nm_rad, moments))
emittance_t = list(map(lambda step_val: step_val[1][5] * nm_rad, moments))
# print(sigx, sigt)
# print beam transversal size over steps
f = plt.figure(figsize=(9, 4.8))
f, axs = plt.subplots(
2, 1, figsize=(9, 4.8), sharex=True, gridspec_kw={"height_ratios": [1, 2]}
)
ax0 = axs[0]
im_xz = ax0.plot(z, x, "--", lw=3, label=r"$x$")
ax0.legend(loc="upper right")
ax0.set_ylim([0, None])
ax0.set_ylabel(r"$x$ [m]")
ax1 = axs[1]
im_sigx = ax1.plot(z, sigx, label=r"$\sigma_x$")
im_sigt = ax1.plot(z, sigt, label=r"$\sigma_t$")
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_t = ax2.plot(z, 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,y}$ [mm-mrad]")
ax2.set_ylabel(r"$\epsilon_{x,t}$ [nm]")
ax1.set_ylim([0, None])
ax2.set_ylim([0, None])
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
plt.savefig("chicane_sigma.png")
else:
plt.show()
# beam transversal scatter plot over steps
plot_ever_nstep = 25 # this is lattice.nslice in inputs
# entry and exit dipedge are part of each bending element
# lattice.elements = sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 sbend2
# dipedge2 drift1 dipedge1 sbend1 drift3
# lattice.nslice = 25
print_steps = [0, 26, 51, 77, 102, 128, 153, 179, 204]
num_plots_per_row = len(print_steps)
fig, axs = plt.subplots(
4, num_plots_per_row, figsize=(9, 6.4), sharex="row", sharey="row"
)
ncol_ax = -1
for step in print_steps:
# plot initial distribution & at exit of each element
print(step)
ncol_ax += 1
# t-pt
ax = axs[(0, ncol_ax)]
beam_at_step = beam[beam["step"] == step]
ax.scatter(
beam_at_step.t.multiply(millimeter), beam_at_step.pt.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[step]:.1f}${z_unit}")
# 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]")
# t-x
ax = axs[(2, ncol_ax)]
beam_at_step = beam[beam["step"] == step]
ax.scatter(
beam_at_step.t.multiply(millimeter), beam_at_step.x.multiply(millimeter), s=0.01
)
ax.set_xlabel(r"$ct$ [mm]")
# t-px
ax = axs[(3, ncol_ax)]
beam_at_step = beam[beam["step"] == step]
ax.scatter(
beam_at_step.t.multiply(millimeter), beam_at_step.px.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