Optimized Triplet
Optimization of focusing parameters for a quadrupole triplet. A 2 GeV electron beam is strongly focused from lattice initial parameters:
\[\begin{split}\beta_x = \beta_y = 40\,\mathrm{m}\\
\alpha_x = -\alpha_y = 2\end{split}\]
to final lattice parameters:
\[\begin{split}\beta_x = \beta_y = 0.55\,\mathrm{m}\\
\alpha_x = \alpha_y = 0\end{split}\]
resulting in a reduction by a factor of 8.5 in the horizontal and vertical beam sizes.
Here, we start with a desired spatial layout of the triplet and find the quadrupole strengths through numerical optimization (by minimizing the L2 norm of alpha and beta) over multiple ImpactX simulations.
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 uses scipy.optimize (methods: Nelder-Mead or L-BFGS-B) to find the quadrupole strengths by minimizing the objective. Conventional optimization algorithms like this work best if there is only a global minima in the objective.
This example can be run as a Python script: python3 run_triplet.py.
#!/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
import impactx
import numpy as np
from impactx import ImpactX, distribution, elements
from scipy.optimize import minimize
# Call MPI_Init and MPI_Finalize only once:
if impactx.Config.have_mpi:
from mpi4py import MPI # noqa
verbose = False
def build_lattice(parameters: tuple, write_particles: bool) -> list:
"""
Create the quadrupole triplet.
Parameters
----------
parameters: tuple
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
Returns
-------
A lattice for ImpactX: a list of impactx.elements.
"""
q1_k, q2_k = parameters
ns = 10 # number of slices per ds in the element
# enforce a mirror symmetry of the triplet
line = [
elements.Drift(ds=2.7, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.2, k=q2_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=2.7, nslice=ns),
]
if write_particles:
monitor = elements.BeamMonitor("monitor", backend="h5")
line = [monitor] + line + [monitor]
return line
def run(parameters: tuple, write_particles=False, write_reduced=False) -> dict:
"""
Run an ImpactX simulation with a new set of lattice parameters.
Parameters
----------
parameters: tuple
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
write_reduced: bool
write the reduced diagnositcs of ImpactX to a file.
Returns
-------
A dictionary with reduced diagnositcs of ImpactX, characterizing
the beam at the end of the simulation.
"""
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("verbose", 0)
sim = ImpactX()
if verbose is False:
sim.verbose = 0
# set numerical parameters and IO control
sim.particle_shape = 2 # B-spline order
sim.space_charge = False
sim.diagnostics = write_reduced
sim.slice_step_diagnostics = write_reduced
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 5 nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 100.0e-12 # 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.0e-4,
lambdaY=2.0e-4,
lambdaT=3.1622776602e-5,
lambdaPx=1.1180339887e-5,
lambdaPy=1.1180339887e-5,
lambdaPt=3.1622776602e-5,
muxpx=0.894427190999916,
muypy=-0.894427190999916,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
sim.lattice.extend(build_lattice(parameters, write_particles=write_particles))
# run simulation
sim.evolve()
# in situ calculate the reduced beam characteristics
beam = sim.particle_container()
rbc = beam.reduced_beam_characteristics()
# clean shutdown
sim.finalize()
return rbc
def objective(parameters: tuple) -> float:
"""
A function that is evaluated by the optimizer.
Parameters
----------
parameters: tuple
quadrupole strengths k of quad 1/3 and quad 2.
Returns
-------
The L2 norm of alpha and beta of the beam at the end of the
simulation.
"""
if verbose:
print(f"Run objective with parameters={parameters}...")
rbc = run(parameters, write_particles=False, write_reduced=False)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
if verbose:
print(f"alpha_x={alpha_x}, alpha_y={alpha_y}, beta_x={beta_x}, beta_y={beta_y}")
alpha_beta_is = np.array([alpha_x, alpha_y, beta_x, beta_y])
beta_x_goal = 0.55
beta_y_goal = beta_x_goal
alpha_beta_goal = np.array([0, 0, beta_x_goal, beta_y_goal])
error = np.sum((alpha_beta_is - alpha_beta_goal) ** 2)
if np.isnan(error):
error = 1.0e99
return error
if __name__ == "__main__":
# Initial guess for the quadrople strengths
initial_quad_strengths = np.array([-3, 3])
# Bounds for values to test: (min, max)
positive = (0, None)
negative = (None, 0)
bounds = [negative, positive]
# optimizer specific values
# https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
# https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html
options = {
"maxiter": 1000,
}
# Call the optimizer
res = minimize(
objective,
initial_quad_strengths,
method="Nelder-Mead",
# method="L-BFGS-B",
tol=1.0e-8,
options=options,
bounds=bounds,
)
# Print the optimization result
print("Optimal parameters for k:", res.x)
print("L2 norm of alpha & beta at the optimum:", res.fun)
# analytical result:
# k: -3.5, 2.75
# alpha & beta: 0, 0, 0.55, 0.55
# final run
rbc = run(res.x, write_particles=True, write_reduced=True)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
print(f"alpha_x={alpha_x} alpha_y={alpha_y}\n beta_x={beta_x} beta_y={beta_y}")
This example uses Xopt (methods: Nelder-Mead or TuRBO) to find the quadrupole strengths by minimizing the objective.
Conventional optimization algorithms like Nelder-Mead work best if there is only a global minima in the objective.
Machine-learning based, surrogate optimization works well for highly dimensional inputs and/or to find global minima in an objective that has potentially many local minima, where conventional optimizers can get stuck.
At the same time, the ML method Bayesian Optimization (BO) is prone to over-explore an objective (at the cost of finding a point closer to the global minima).
The variation Trust Region Bayesian Optimization (TuRBO) was developed to narrow down further on found minima.
This example can be run as a Python script: python3 tests/python/test_xopt.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 The ImpactX Community
#
# Authors: Axel Huebl
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import importlib
import amrex.space3d as amr
import impactx
import numpy as np
import pytest
from impactx import ImpactX, distribution, elements
# configure the test
verbose = True
gen_name = "Nelder-Mead" # TuRBO or Nelder-Mead
max_steps = 60
def build_lattice(parameters: dict, write_particles: bool) -> list:
"""
Create the quadrupole triplet.
Parameters
----------
parameters: dict
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
Returns
-------
A lattice for ImpactX: a list of impactx.elements.
"""
q1_k, q2_k = parameters["q1_k"], parameters["q2_k"]
ns = 10 # number of slices per ds in the element
# enforce a mirror symmetry of the triplet
line = [
elements.Drift(ds=2.7, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.2, k=q2_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=2.7, nslice=ns),
]
if write_particles:
monitor = elements.BeamMonitor("monitor", backend="h5")
line = [monitor] + line + [monitor]
return line
def run(parameters: dict, write_particles=False, write_reduced=False) -> dict:
"""
Run an ImpactX simulation with a new set of lattice parameters.
Parameters
----------
parameters: dict
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
write_reduced: bool
write the reduced diagnositcs of ImpactX to a file.
Returns
-------
A dictionary with reduced diagnositcs of ImpactX, characterizing
the beam at the end of the simulation.
"""
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("verbose", 0)
sim = ImpactX()
sim.verbose = 0
# set numerical parameters and IO control
sim.particle_shape = 2 # B-spline order
sim.space_charge = False
sim.diagnostics = write_reduced
sim.slice_step_diagnostics = write_reduced
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 5 nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 100.0e-12 # 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.0e-4,
lambdaY=2.0e-4,
lambdaT=3.1622776602e-5,
lambdaPx=1.1180339887e-5,
lambdaPy=1.1180339887e-5,
lambdaPt=3.1622776602e-5,
muxpx=0.894427190999916,
muypy=-0.894427190999916,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
sim.lattice.extend(build_lattice(parameters, write_particles=write_particles))
# run simulation
sim.evolve()
# in situ calculate the reduced beam characteristics
beam = sim.particle_container()
rbc = beam.reduced_beam_characteristics()
# clean shutdown
sim.finalize()
return rbc
def objective(parameters: dict) -> dict:
"""
A function that is evaluated by the optimizer.
Parameters
----------
parameters: dict
quadrupole strengths k of quad 1/3 and quad 2.
Returns
-------
The L2 norm of alpha and beta of the beam at the end of the
simulation.
"""
if verbose:
print(f"Run objective with parameters={parameters}...")
rbc = run(parameters, write_particles=False, write_reduced=False)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
if verbose:
print(
f" -> alpha_x={alpha_x}, alpha_y={alpha_y}, beta_x={beta_x}, beta_y={beta_y}"
)
alpha_beta_is = np.array([alpha_x, alpha_y, beta_x, beta_y])
beta_x_goal = 0.55
beta_y_goal = beta_x_goal
alpha_beta_goal = np.array([0, 0, beta_x_goal, beta_y_goal])
error = np.sum((alpha_beta_is - alpha_beta_goal) ** 2)
# xopt will ignore NaN results, no cleaning needed
rbc["error"] = error
return rbc
@pytest.mark.skipif(
importlib.util.find_spec("xopt") is None, reason="xopt is not available"
)
def test_xopt():
from xopt import Xopt
from xopt.evaluator import Evaluator
from xopt.vocs import VOCS
if gen_name == "TuRBO":
from xopt.generators.bayesian import UpperConfidenceBoundGenerator as Generator
gen_args = {"turbo_controller": "optimize"}
elif gen_name == "Nelder-Mead":
from xopt.generators.scipy.neldermead import NelderMeadGenerator as Generator
gen_args = {}
else:
raise RuntimeError(f"Unconfigured generator named '{gen_name}'")
# Bounds for values to test: (min, max)
positive = [0, 10]
negative = [-10, 0]
# define variables and function objectives
vocs = VOCS(
variables={
"q1_k": negative,
"q2_k": positive,
},
objectives={"error": "MINIMIZE"},
)
# create Xopt evaluator, generator, and Xopt objects
evaluator = Evaluator(function=objective)
generator = Generator(vocs=vocs, **gen_args)
X = Xopt(evaluator=evaluator, generator=generator, vocs=vocs)
# Initial guess for the quadrople strengths
initial_quad_strengths = {
"q1_k": np.array([-3]),
"q2_k": np.array([3]),
}
if gen_name == "TuRBO":
# a few random guesses
# X.random_evaluate(3)
# a few somewhat educated guesses
X.evaluate_data(initial_quad_strengths)
elif gen_name == "Nelder-Mead":
# a few somewhat educated guesses
X.generator.initial_point = initial_quad_strengths
# run optimization for 60 steps (iterations)
for i in range(max_steps):
X.step()
# Print all trials
if verbose:
print(X.data)
# plot
# X.vocs.normalize_inputs(X.data).plot(*X.vocs.variable_names, kind="scatter")
# Select the best result
best_idx, best_error = X.vocs.select_best(X.data)
best_run = X.data.iloc[best_idx]
best_ks = best_run[["q1_k", "q2_k"]].to_dict(orient="index")[best_idx[0]]
# Print the optimization result
print("Optimal parameters for k:", best_ks)
print("L2 norm of alpha & beta at the optimum:", best_run["error"].values[0])
# analytical result:
# k: -3.5, 2.75
# alpha & beta: 0, 0, 0.55, 0.55
# final run w/ detailed I/O on
rbc = run(best_ks, write_particles=True, write_reduced=True)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
print(f"alpha_x={alpha_x} alpha_y={alpha_y}\n beta_x={beta_x} beta_y={beta_y}")
if __name__ == "__main__":
# Call MPI_Init and MPI_Finalize only once:
if impactx.Config.have_mpi:
from mpi4py import MPI # noqa
test_xopt()
Analyze
We run the following script to analyze correctness:
Script analysis_triplet.py
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
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.0 * 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],
[
4.47213595500e-004,
4.47213595500e-004,
3.16227766017e-005,
5.0e-009,
5.0e-009,
1.0e-009,
],
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.0 * 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.2440442742e-005,
5.2440442742e-005,
3.16227766017e-005,
5.0e-009,
5.0e-009,
1.0e-009,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the optimized beam evolution over time:
Script plot_triplet.py
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import argparse
import os
import matplotlib.pyplot as plt
import numpy as np
# options to run this script
parser = argparse.ArgumentParser(description="Plot the quadrupole triplet benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()
# read reduced diagnostics
rdc_name = "diags/reduced_beam_characteristics.0"
if os.path.exists(rdc_name):
data = np.loadtxt(rdc_name, skiprows=1)
else: # OpenMP
data = np.loadtxt(rdc_name + ".0", skiprows=1)
s = data[:, 1]
beta_x = data[:, 33]
beta_y = data[:, 34]
xMin = 0.0
xMax = 8.6
yMin = 0.0
yMax = 65.0
# Plotting
plt.figure(figsize=(10, 6))
plt.xscale("linear")
plt.yscale("linear")
plt.xlim([xMin, xMax])
# plt.ylim([yMin, yMax])
plt.xlabel("s (m)", fontsize=30)
plt.ylabel("CS Twiss beta (m)", fontsize=30)
plt.xticks(fontsize=25)
plt.yticks(fontsize=25)
plt.grid(True)
# Plot the data
plt.plot(s, beta_x, "b", label="Horizontal", linewidth=2, linestyle="solid")
plt.plot(s, beta_y, "r", label="Vertical", linewidth=2, linestyle="solid")
# Show plot
plt.legend(fontsize=20)
plt.tight_layout()
if args.save_png:
plt.savefig("triplet.png")
else:
plt.show()
Fig. 7 CS Twiss beta