Turn-Dependend Element Changes
10 periods of drift elements, where each drift is a stand-in for a “turn” map. On each turn, the next turn map (drift) is modified based in particle properties.
A 2 GeV electron bunch with normalized transverse rms emittance of 50 pm.
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:
Python script:
python3 run_turn_update.pyor
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: Chad Mitchell, Axel Huebl
# 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.slice_step_diagnostics = True
sim.tiny_profiler = False
# 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 = 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_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=5.0e-6, # 5 um
lambdaY=8.0e-6, # 8 um
lambdaT=0.0599584916, # 200 ps
lambdaPx=2.5543422003e-9, # exn = 50 pm-rad
lambdaPy=1.5964638752e-9, # eyn = 50 pm-rad
lambdaPt=9.0e-4, # approximately dE/E
muxpx=0.0,
muypy=0.0,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
ns = 1 # number of slices per ds in the element
bend = [
monitor,
elements.Drift(ds=1.0, nslice=ns, name="drift1"),
]
# assign a lattice segment
sim.lattice.extend(bend)
sim.periods = 10
# at the beginning of each period (e.g, turn or channel period), modify the lattice
def hook_before_period(sim):
turn = sim.tracking_period
step = sim.tracking_step
print("- Before turn:")
print(f" Updating lattice at turn {turn}, step {step}", flush=True)
beam = sim.particle_container()
ref = beam.ref_particle()
rbc = beam.beam_moments()
print(
f" Beam at s={ref.s:.2f}m, t={ref.t:.2f}s with beta_x={rbc['beta_x']}m",
flush=True,
)
if turn > 0:
next_ds = 1.0 / ref.s
print(f" Updating next drift ds to: {next_ds:.2f}m", flush=True)
sim.lattice.select(name="drift1")[0].ds = next_ds
def hook_before_element(sim):
element = sim.tracking_element
print("- Before element:")
print(
f" Current element name: {element.name} with ds={element.ds:.2f}", flush=True
)
if element.name != "monitor":
print(f" Drift ds is: {element.ds:.2f}m", flush=True)
sim.hook["before_period"] = hook_before_period
sim.hook["before_element"] = hook_before_element
# run simulation
sim.track_particles()
# clean shutdown
sim.finalize()
Analyze
We run the following script to analyze correctness:
Script analysis_turn_update.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()
s_final = series.iterations[last_step].particles["beam"].get_attribute("s_ref")
# 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}"
f" s_final={s_final}"
)
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, s_final],
[
5.0e-6,
8.0e-6,
0.0599584916,
1.277171100130637e-14,
1.277171100130637e-14,
5.396264243e-005,
4.339439692724345,
],
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 = 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],
[
5.0e-6,
8.0e-6,
0.0599584916,
1.277171100130637e-14,
1.277171100130637e-14,
5.396264243e-005,
],
rtol=rtol,
atol=atol,
)