Multitask optimization with FBPIC and Wake-T#
Description#
This is an advanced example that shows how perform a multitask Bayesian optimization using two simulations codes of different fidelity (FBPIC and Wake-T). The scripts provided here can be used to reproduce the results from the paper
“Bayesian optimization of laser-plasma accelerators assisted by reduced physical models” by A. Ferran Pousa, S. Jalas, M. Kirchen, A. Martinez de la Ossa, M. Thévenet, J. Larson, S. Hudson, A. Huebl, J.-L. Vay, and R. Lehe (link).
Requirements#
In addition to optimas, the following packages should be installed:
Scripts#
Files included:
run_opt.py: defines and launches the optimization with optimas.template_simulation_script.py: template used by optimas to generate the FBPIC and Wake-T simulation scripts.analysis_script.py: defines how the simulation data is analyzed to yield the value of the objective function.bunch_utils.py: contains methods for generating the beam particle distributions given to the simulations.custom_fld_diags.py: custom FBPIC field diagnostics that have been to generate the output with the same location and periodicity as Wake-T.custom_ptcl_diags.py: custom FBPIC particle diagnostics that have been modified to generate the output with the same location and periodicity as Wake-T.
You can have a look at the main scripts below:
"""Multitask optimization of an LPA with Wake-T and FBPIC."""
from multiprocessing import set_start_method
from optimas.core import VaryingParameter, Objective, Parameter, Task
from optimas.generators import AxMultitaskGenerator
from optimas.evaluators import TemplateEvaluator, MultitaskEvaluator
from optimas.explorations import Exploration
from analysis_script import analyze_simulation
# Create varying parameters and objectives.
var_1 = VaryingParameter("beam_i_1", 1.0, 10.0) # kA
var_2 = VaryingParameter("beam_i_2", 1.0, 10.0) # kA
var_3 = VaryingParameter("beam_z_i_2", -10.0, 10.0) # µm
var_4 = VaryingParameter("beam_length", 1.0, 20.0) # µm
obj = Objective("f", minimize=True)
# Define other quantities to analyze (which are not the optimization objective)
par_1 = Parameter("energy_med")
par_2 = Parameter("energy_mad")
par_3 = Parameter("charge")
# Create tasks.
lofi_task = Task("wake-t", n_init=96, n_opt=96)
hifi_task = Task("fbpic", n_init=3, n_opt=3)
# Create generator.
gen = AxMultitaskGenerator(
varying_parameters=[var_1, var_2, var_3, var_4],
objectives=[obj],
analyzed_parameters=[par_1, par_2, par_3],
use_cuda=True,
dedicated_resources=True,
hifi_task=hifi_task,
lofi_task=lofi_task,
)
# Create evaluators for each task.
ev_lofi = TemplateEvaluator(
sim_template="template_simulation_script.py",
analysis_func=analyze_simulation,
sim_files=["bunch_utils.py", "custom_fld_diags.py", "custom_ptcl_diags.py"],
)
ev_hifi = TemplateEvaluator(
sim_template="template_simulation_script.py",
analysis_func=analyze_simulation,
sim_files=["bunch_utils.py", "custom_fld_diags.py", "custom_ptcl_diags.py"],
)
# Create a multitask evaluator. This associates each task to each task
# evaluator.
ev = MultitaskEvaluator(
tasks=[lofi_task, hifi_task], task_evaluators=[ev_lofi, ev_hifi]
)
# Create exploration.
n_batches = 50
exp = Exploration(
generator=gen,
evaluator=ev,
max_evals=(
(lofi_task.n_opt + hifi_task.n_opt) * n_batches
+ lofi_task.n_init
+ hifi_task.n_init
),
sim_workers=96,
run_async=False,
)
# Run exploration.
if __name__ == "__main__":
set_start_method("spawn")
exp.run()
"""Script for simulating an LPA with external injection with Wake-T and FBPIC.
The simulation code is determined from the `task` parameter. The beam
current, position and length are parameters exposed to the optimizer to
try to achieve optimal beam loading.
"""
import numpy as np
import scipy.constants as ct
from wake_t import GaussianPulse, PlasmaStage, ParticleBunch
import aptools.plasma_accel.general_equations as ge
from bunch_utils import trapezoidal_bunch
# Parammeters exposed to optimizer.
task = {{task}}
beam_i_1 = {{beam_i_1}}
beam_i_2 = {{beam_i_2}}
beam_z_i_2 = {{beam_z_i_2}}
beam_length = {{beam_length}}
# General simulation parameters.
n_p_plateau = 2e23
l_plateau = 10e-2
w0_laser = 40e-6
z_beam = (50 + beam_z_i_2) * 1e-6
l_beam = beam_length * 1e-6
i1_beam = beam_i_1 * 1e3
i2_beam = beam_i_2 * 1e3
def run_simulation():
"""Run a simulation of the LPA with Wake-T or FBPIC."""
# Base laser parameters.
E_laser = 10 # J
tau_laser = 25e-15 # s (fwhm)
lambda0 = 0.8e-6 # m
a0 = determine_laser_a0(E_laser, tau_laser, w0_laser, lambda0)
# Base beam parameters.
E_beam = 200 # MeV
gamma_beam = E_beam / 0.511
n_emitt_x = 3e-6
n_emitt_y = 0.5e-6
ene_sp = 0.1 # %
n_part = 1e5
sz0 = 1e-6 # gaussian decay
kp = np.sqrt(n_p_plateau * ct.e**2 / (ct.epsilon_0 * ct.m_e * ct.c**2))
kbeta = kp / np.sqrt(2.0 * gamma_beam) # betatron wavenumber (blowout)
betax0 = 1.0 / kbeta # matched beta
sx0 = np.sqrt(n_emitt_x * betax0 / gamma_beam) # matched beam size (rms)
sy0 = np.sqrt(n_emitt_y * betax0 / gamma_beam) # matched beam size (rms)
# Determine guiding channel.
r_e = ct.e**2 / (4.0 * np.pi * ct.epsilon_0 * ct.m_e * ct.c**2)
rel_delta_n_over_w2 = 1.0 / (np.pi * r_e * w0_laser**4 * n_p_plateau)
# Generate bunch
x, y, z, ux, uy, uz, q = trapezoidal_bunch(
i1_beam,
i2_beam,
n_part=n_part,
gamma0=gamma_beam,
s_g=ene_sp * gamma_beam / 100,
length=l_beam,
s_z=sz0,
emit_x=n_emitt_x,
s_x=sx0,
emit_y=n_emitt_y,
s_y=sy0,
zf=0.0,
tf=0.0,
)
z -= l_beam / 2 + z_beam
w = np.abs(q / ct.e)
bunch = ParticleBunch(w, x, y, z, ux, uy, uz, name="bunch")
# Distance between right boundary and laser centroid.
dz_lb = 4.0 * ct.c * tau_laser
# Maximum radial extension of the plasma.
p_rmax = 2.5 * w0_laser
# Box length.
l_box = dz_lb + 90e-6
# Number of diagnostics
n_out = 3
if task == "fbpic":
run_fbpic(
a0,
w0_laser,
tau_laser,
lambda0,
bunch,
n_p_plateau,
l_plateau,
rel_delta_n_over_w2,
p_rmax,
dz_lb,
l_box,
n_out,
)
elif task == "wake-t":
run_wake_t(
a0,
w0_laser,
tau_laser,
lambda0,
bunch,
n_p_plateau,
l_plateau,
rel_delta_n_over_w2,
p_rmax,
dz_lb,
l_box,
n_out - 1,
)
def determine_laser_a0(ene, tau_fwhm, w0, lambda0):
"""Determine the laser a0 from the energy, size and wavelength."""
tau = tau_fwhm / np.sqrt(2.0 * np.log(2))
k0 = 2.0 * np.pi / lambda0 # Laser wavenumber
PA = ct.epsilon_0 * ct.c**5 * ct.m_e**2 / ct.e**2 # Power constant
P0 = ene / (np.sqrt(2 * np.pi) * (tau / 2))
i0 = P0 / ((np.pi / 2) * w0**2)
a0 = np.sqrt(i0 / (PA * k0**2 / 2))
return a0
def density_profile(z):
"""Define the longitudinal density profile of the plasma."""
# Allocate relative density
n = np.ones_like(z)
# Make zero before plateau
n = np.where(z < 0, 0, n)
# Make zero after plateau
n = np.where(z >= l_plateau, 0, n)
# Return absolute density
return n * n_p_plateau
def run_wake_t(
a0,
w0,
tau_fwhm,
lambda0,
bunch,
n_p,
l_plasma,
pc,
p_rmax,
dz_lb,
l_box,
n_out,
):
"""Run a Wake-T simulation of the LPA."""
# Create laser.
laser = GaussianPulse(
xi_c=0.0, l_0=lambda0, w_0=w0, a_0=a0, tau=tau_fwhm, z_foc=0.0
)
# Plasma stage.
s_d = ge.plasma_skin_depth(n_p * 1e-6)
dr = s_d / 20
dz = tau_fwhm * ct.c / 40
r_max = w0 * 4
plasma = PlasmaStage(
l_plasma,
density=density_profile,
wakefield_model="quasistatic_2d",
n_out=n_out,
laser=laser,
laser_evolution=True,
r_max=r_max,
r_max_plasma=p_rmax,
xi_min=dz_lb - l_box,
xi_max=dz_lb,
n_r=int(r_max / dr),
n_xi=int(l_box / dz),
dz_fields=l_box * 2,
ppc=4,
parabolic_coefficient=pc,
max_gamma=25,
dt_bunch=calculate_waket_timestep(bunch, n_p),
bunch_pusher="boris",
)
# Do tracking.
plasma.track(bunch, opmd_diag=True, diag_dir="diags")
def run_fbpic(
a0,
w0,
tau_fwhm,
lambda0,
bunch,
n_p,
l_plasma,
pc,
p_rmax,
dz_lb,
l_box,
n_out,
):
"""Run an FBPIC simulation of the LPA."""
from fbpic.main import Simulation
from fbpic.lpa_utils.boosted_frame import BoostConverter
from fbpic.lpa_utils.bunch import add_particle_bunch_from_arrays
from fbpic.lpa_utils.laser import add_laser
from custom_ptcl_diags import BackTransformedParticleDiagnostic
use_cuda = True
n_order = -1
# Boosted frame
gamma_boost = 25.0
boost = BoostConverter(gamma_boost)
# The laser (Gaussian)
lambda0 = 0.8e-6 # Laser wavelength
# Laser duration (2 sigmas) in intensity
tau = tau_fwhm / np.sqrt(2.0 * np.log(2))
ctau = tau * ct.c
# The simulation box
zmin = -l_box # Left edge of the simulation box (meters)
zmax = 0.0e-6 # Right edge of the simulation box (meters)
rmax = w0 * 4
dz_adv = lambda0 / 80.0 # Advised longitudinal resolution
Nz_adv = int(l_box / dz_adv)
Nz = Nz_adv # Number of gridpoints along z
Nm = 3 # Number of modes used
s_d = ge.plasma_skin_depth(n_p * 1e-6)
dr = s_d / 20
Nr = int(rmax / dr)
# Laser centroid
z0 = zmax - dz_lb
# The simulation timestep
dz = (zmax - zmin) / Nz
dt = dz / ct.c
# The moving window
v_window = ct.c # velocity of the window
# Velocity of the Galilean frame (for suppression of the NCI)
(v_comoving,) = boost.velocity([0.0])
# ------------
# The plasma particles
p_zmin = zmax # Position of the beginning of the plasma (meters)
p_nz = 2 # Number of particles per cell along z
p_nr = 2 # Number of particles per cell along r
p_nt = 8 # Number of particles per cell along theta
# The interaction length of the simulation (meters)
L_lab_interact = l_plasma
# Duration of plasma interaction (i.e. the time it takes for the moving
# window to slide across the plasma)
T_lab_interact_plasma = L_lab_interact / v_window
# Number of discrete diagnostic snapshots in the lab frame
N_lab_diag = n_out
# data dumping period in dt units:
dt_lab_diag_period = T_lab_interact_plasma / (
N_lab_diag - 1
) # Period of the diagnostics (seconds)
# In boosted frame:
(v_window_boosted,) = boost.velocity([v_window])
# Interaction time in boosted frame
T_interact = boost.interaction_time(L_lab_interact, (zmax - zmin), v_window)
# Period of writing the cached backtransformed lab frame diagnostics to
# disk (in number of iterations)
write_period = 200
# Density function
def dens_func(z, r):
z_lab = z * gamma_boost
n = density_profile(z_lab) / n_p
n = n * (1.0 + pc * r**2)
return n
# External bunch
if bunch is not None:
x, y, z, px, py, pz, w = (
bunch.x,
bunch.y,
bunch.xi,
bunch.px,
bunch.py,
bunch.pz,
bunch.w,
)
z += z0
# Initialize the simulation object
sim = Simulation(
Nz=Nz,
zmax=zmax,
Nr=Nr,
rmax=rmax,
Nm=Nm,
dt=dt,
zmin=zmin,
v_comoving=v_comoving,
gamma_boost=boost.gamma0,
n_order=n_order,
use_cuda=use_cuda,
boundaries={"z": "open", "r": "reflective"},
particle_shape="cubic",
)
# Add the Helium ions (full pre-ionized: levels 1 and 2)
sim.add_new_species(
q=ct.e,
m=ct.m_p,
n=n_p,
dens_func=dens_func,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=p_zmin,
p_rmax=p_rmax,
)
# Plasma electrons: coming from helium
sim.add_new_species(
q=-ct.e,
m=ct.m_e,
n=n_p,
dens_func=dens_func,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=p_zmin,
p_rmax=p_rmax,
)
# Add an electron bunch
if bunch is not None:
add_particle_bunch_from_arrays(
sim=sim,
q=-ct.e,
m=ct.m_e,
x=x,
y=y,
z=z,
ux=px,
uy=py,
uz=pz,
w=w,
boost=boost,
z_injection_plane=0.0,
)
# Add a laser to the fields of the simulation
add_laser(
sim=sim,
a0=a0,
w0=w0,
ctau=ctau,
z0=z0,
lambda0=lambda0,
zf=0.0,
gamma_boost=boost.gamma0,
method="antenna",
z0_antenna=0.0,
cep_phase=np.pi,
)
# Configure the moving window
sim.set_moving_window(v=v_window_boosted)
# Add diagnostics
write_dir = "diags"
# Set start time of diagnostics to the exact moment the bunch enters the
# plasma. (Required to have output at same location as Wake-T)
if bunch is not None:
T_start_lab = (zmax - np.average(z)) / v_window
else:
T_start_lab = 0.0
# Add diagnostics.
sim.diags = []
# sim.diags = [
# BackTransformedFieldDiagnostic(
# zmin,
# zmax,
# v_window,
# T_start_lab,
# dt_lab_diag_period,
# N_lab_diag,
# boost.gamma0,
# fieldtypes=["E", "B", "rho"],
# period=write_period,
# fldobject=sim.fld,
# comm=sim.comm,
# write_dir=write_dir,
# )
# ]
if bunch is not None:
sim.diags += [
BackTransformedParticleDiagnostic(
zmin_lab=zmin,
zmax_lab=zmax,
v_lab=v_window,
t_start_lab=T_start_lab,
dt_snapshots_lab=dt_lab_diag_period,
Ntot_snapshots_lab=N_lab_diag,
gamma_boost=boost.gamma0,
period=write_period,
fldobject=sim.fld,
species={"bunch": sim.ptcl[2]},
comm=sim.comm,
write_dir=write_dir,
)
]
# Number of iterations to perform
N_step = int(T_interact / sim.dt) + write_period
# Run the simulation
sim.step(N_step)
print("")
def calculate_waket_timestep(beam, n_p):
"""Calculate the timestep of the bunch pusher in Wake-T."""
mean_gamma = np.sqrt(np.average(beam.pz) ** 2 + 1)
# calculate maximum focusing along stage.
w_p = np.sqrt(n_p * ct.e**2 / (ct.m_e * ct.epsilon_0))
max_kx = (ct.m_e / (2 * ct.e * ct.c)) * w_p**2
w_x = np.sqrt(ct.e * ct.c / ct.m_e * max_kx / mean_gamma)
period_x = 1 / w_x
dt = 0.1 * period_x
return dt
if __name__ == "__main__":
run_simulation()
"""Defines the analysis function that runs after the simulation."""
import os
import numpy as np
import matplotlib.pyplot as plt
import visualpic as vp
from aptools.plotting.quick_diagnostics import (
phase_space_overview,
slice_analysis,
)
def analyze_simulation(simulation_directory, output_params):
"""Analyze the output of the simulation."""
# Load data.
diags_dir = os.path.join(simulation_directory, "diags/hdf5")
dc = vp.DataContainer("openpmd", diags_dir)
dc.load_data()
# Get final bunch distribution.
bunch = dc.get_species("bunch")
ts = bunch.timesteps
bunch_data = bunch.get_data(ts[-1])
x = bunch_data["x"][0]
y = bunch_data["y"][0]
z = bunch_data["z"][0]
px = bunch_data["px"][0]
py = bunch_data["py"][0]
pz = bunch_data["pz"][0]
q = bunch_data["q"][0]
# Remove particles with pz < 100
pz_filter = np.where(pz >= 100)
x = x[pz_filter]
y = y[pz_filter]
z = z[pz_filter]
px = px[pz_filter]
py = py[pz_filter]
pz = pz[pz_filter]
q = q[pz_filter]
# Calculate relevant parameters.
q_tot = np.abs(np.sum(q)) * 1e12 # pC
q_ref = 10 # pC
# ene = np.average(pz, weights=q) * 0.511 # MeV
med, mad = weighted_mad(pz * 0.511, q)
mad_rel = mad / med
med *= 1e-3 # GeV
mad_rel_ref = 1e-2
# Calculate objective.
f = np.log(med * q_tot / q_ref / (mad_rel / mad_rel_ref))
# Store quantities in output.
output_params["f"] = -f
output_params["charge"] = q_tot
output_params["energy_med"] = med
output_params["energy_mad"] = mad
# Save objective to file (for convenience).
np.savetxt("f.txt", np.array([f]))
# Make plots.
try:
plt.figure()
slice_analysis(x, y, z, px, py, pz, q, show=False)
plt.savefig("final_lon_phase_space.png")
plt.figure()
phase_space_overview(x, y, z, px, py, pz, q, show=False)
plt.savefig("final_phase_space.png")
except Exception:
print("Failed to make plots.")
# Remove all diagnostics except last file.
try:
for file in sorted(os.listdir(diags_dir))[:-1]:
file_path = os.path.join(diags_dir, file)
os.remove(file_path)
except Exception:
print("Could not remove diagnostics.")
return output_params
def weighted_mad(x, w):
"""Calculate weighted median absolute deviation."""
med = weighted_median(x, w)
mad = weighted_median(np.abs(x - med), w)
return med, mad
def weighted_median(data, weights):
"""Compute the weighted quantile of a 1D numpy array.
Parameters
----------
data : ndarray
Input array (one dimension).
weights : ndarray
Array with the weights of the same size of `data`.
quantile : float
Quantile to compute. It must have a value between 0 and 1.
Returns
-------
quantile_1D : float
The output value.
"""
quantile = 0.5
# Check the data
if not isinstance(data, np.matrix):
data = np.asarray(data)
if not isinstance(weights, np.matrix):
weights = np.asarray(weights)
nd = data.ndim
if nd != 1:
raise TypeError("data must be a one dimensional array")
ndw = weights.ndim
if ndw != 1:
raise TypeError("weights must be a one dimensional array")
if data.shape != weights.shape:
raise TypeError("the length of data and weights must be the same")
if (quantile > 1.0) or (quantile < 0.0):
raise ValueError("quantile must have a value between 0. and 1.")
# Sort the data
ind_sorted = np.argsort(data)
sorted_data = data[ind_sorted]
sorted_weights = weights[ind_sorted]
# Compute the auxiliary arrays
Sn = np.cumsum(sorted_weights)
# TODO: Check that the weights do not sum zero
# assert Sn != 0, "The sum of the weights must not be zero"
Pn = (Sn - 0.5 * sorted_weights) / Sn[-1]
# Get the value of the weighted median
return np.interp(quantile, Pn, sorted_data)