Aperture Collimation
Proton beam undergoing collimation by a rectangular boundary aperture.
We use a 250 MeV proton beam with a horizontal rms beam size of 1.56 mm and a vertical rms beam size of 2.21 mm.
After a short drift of 0.123, the beam is scraped by a 1 mm x 1.5 mm rectangular aperture.
In this test, the initial values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\epsilon_t\) must agree with nominal values. The test fails if:
any of the final coordinates for the valid (not lost) particles lie outside the aperture boundary or
any of the lost particles are inside the aperture boundary or
if the sum of lost and kept particles is not equal to the initial particles or
if the recorded position \(s\) for the lost particles does not coincide with the drift distance.
Run
This example can be run as a Python script (python3 run_aperture.py) or with an app with an input file (impactx input_aperture.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.space3d as amr
from impactx import ImpactX, distribution, elements
# work-around for https://github.com/ECP-WarpX/impactx/issues/499
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("the_arena_is_managed", 1)
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
sim.particle_lost_diagnostics_backend = "h5"
# domain decomposition & space charge mesh
sim.init_grids()
# load a 250 MeV proton beam with an initial
# horizontal rms emittance of 1 um and an
# initial vertical rms emittance of 2 um
kin_energy_MeV = 250.0 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.559531175539e-3,
lambdaY=2.205510139392e-3,
lambdaT=1.0e-3,
lambdaPx=6.41218345413e-4,
lambdaPy=9.06819680526e-4,
lambdaPt=1.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.extend(
[
monitor,
elements.Drift(0.123),
elements.Aperture(xmax=1.0e-3, ymax=1.5e-3, shape="rectangular"),
monitor,
]
)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.559531175539e-3
beam.lambdaY = 2.205510139392e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 6.41218345413e-4
beam.lambdaPy = 9.06819680526e-4
beam.lambdaPt = 1.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift collimator monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
drift.type = drift
drift.ds = 0.123
collimator.type = aperture
collimator.shape = rectangular
collimator.xmax = 1.0e-3
collimator.ymax = 1.5e-3
# work-around for https://github.com/ECP-WarpX/impactx/issues/499
amrex.the_arena_is_managed = 1
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
diag.backend = h5
Analyze
We run the following script to analyze correctness:
Script analysis_aperture.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()
series_lost = io.Series("diags/openPMD/particles_lost.h5", io.Access.read_only)
particles_lost = series_lost.iterations[0].particles["beam"].to_df()
# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
# we lost particles in apertures
assert num_particles > len(final)
assert num_particles == len(particles_lost) + 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 = 1.8 * 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.559531175539e-3,
2.205510139392e-3,
1.0e-3,
1.0e-6,
2.0e-6,
1.0e-6,
],
rtol=rtol,
atol=atol,
)
# particle-wise comparison against the rectangular aperture boundary
xmax = 1.0e-3
ymax = 1.5e-3
# kept particles
dx = abs(final["position_x"]) - xmax
dy = abs(final["position_y"]) - ymax
print()
print(f" x_max={final['position_x'].max()}")
print(f" x_min={final['position_x'].min()}")
assert np.less_equal(dx.max(), 0.0)
print(f" y_max={final['position_y'].max()}")
print(f" y_min={final['position_y'].min()}")
assert np.less_equal(dy.max(), 0.0)
# lost particles
dx = abs(particles_lost["position_x"]) - xmax
dy = abs(particles_lost["position_y"]) - ymax
print()
print(f" x_max={particles_lost['position_x'].max()}")
print(f" x_min={particles_lost['position_x'].min()}")
assert np.greater_equal(dx.max(), 0.0)
print(f" y_max={particles_lost['position_y'].max()}")
print(f" y_min={particles_lost['position_y'].min()}")
assert np.greater_equal(dy.max(), 0.0)
# check that s is set correctly
lost_at_s = particles_lost["s_lost"]
drift_s = np.ones_like(lost_at_s) * 0.123
assert np.allclose(lost_at_s, drift_s)