Optimization with HiPACE++#
Description#
This examples shows how to perform a Bayesian optimization of a PWFA using HiPACE++.
The setup is a simple driver-witness configuration where the witness is optimized to maximize the objetive
\[f = \frac{\sqrt{Q} E_{MED}}{100 E_{MAD}}\]
where \(Q\) is the beam charge, \(E_{MED}\) is the median energy, and \(E_{MAD}\) is the median absolute deviation energy spread. The only optimization parameter is the charge:
'witness_charge': parameter in the range \([0.05, 1.]\) in units of \(\mathrm{nC}\).
The optimization is carried out using an
AxSingleFidelityGenerator and a
TemplateEvaluator. In this case, the function
analyze_simulation that analyzes the output of each simulation is defined
in a separate file analysis_script.py and imported into the main
optimas script.
The example is set up to make use of a system of 4 GPUs, where each HiPACE++ simulation uses a 2 GPUs and 2 simulations are carried out in parallel.
If HiPACE++ is installed in a different environment than optimas, make
sure to specify env_script and env_mpi in the TemplateEvaluator.
Scripts#
The files needed to run the optimization should be located in a folder
(named e.g., optimization) with the following structure:
optimization
├── run_example.py
├── template_simulation_script.py
└── analysis_script.py
The optimization is started by executing:
python run_example.py
The scripts needed to run this example can be seen below.
"""Example Bayesian optimization a PWFA with HiPACE++.
The HiPACE++ simulations are performed using the template defined in the
`template_simulation_script` file.
In addition to the objective `f`, three additional parameters
are analyzed for each simulation and including in the optimization
history. The calculation of `f` and the additional parameters is performed
in the `analyze_simulation` function, which for convenience is here defined in
the `analysis_script.py` file.
"""
from optimas.core import Parameter, VaryingParameter, Objective
from optimas.generators import AxSingleFidelityGenerator
from optimas.evaluators import TemplateEvaluator
from optimas.explorations import Exploration
from analysis_script import analyze_simulation
# Create varying parameters and objectives.
var_1 = VaryingParameter("witness_charge", 0.05, 1.0)
obj = Objective("f", minimize=False)
# Define additional parameters to analyze.
energy_med = Parameter("energy_med")
energy_mad = Parameter("energy_mad")
charge = Parameter("charge")
# Create generator.
gen = AxSingleFidelityGenerator(
varying_parameters=[var_1],
objectives=[obj],
analyzed_parameters=[energy_med, energy_mad, charge],
n_init=4,
)
# Create evaluator.
ev = TemplateEvaluator(
sim_template="template_simulation_script",
analysis_func=analyze_simulation,
executable="/path/to/build/bin/hipace",
n_gpus=2, # Use 2 GPUs per simulation.
# Uncomment if HiPACE is installed in a different environment than optimas.
# env_script='/path/to/profile.hipace',
# Uncomment if `env_script` loads a different MPI to that used by optimas.
# env_mpi='openmpi'
)
# Create exploration.
exp = Exploration(generator=gen, evaluator=ev, max_evals=100, sim_workers=2)
# To safely perform exploration, run it in the block below (this is needed
# for some flavours of multiprocessing, namely spawn and forkserver)
if __name__ == "__main__":
exp.run()
max_step = 300
amr.n_cell = 256 256 256
amr.max_level = 0
hipace.max_time = 0.3/clight
diagnostic.output_period = 30
hipace.verbose = 1
hipace.depos_order_xy = 2
hipace.dt = adaptive
hipace.nt_per_betatron = 30
geometry.is_periodic = true true false # Is periodic?
geometry.prob_lo = -250.e-6 -250.e-6 -250.e-6 # physical domain
geometry.prob_hi = 250.e-6 250.e-6 110.e-6
beams.names = driver witness
driver.position_mean = 0. 0. 0.
driver.position_std = 2.e-6 2.e-6 30.e-6
driver.injection_type = fixed_weight
driver.num_particles = 1000000
driver.total_charge = .6e-9
driver.u_mean = 0. 0. 1000.
driver.u_std = 2. 2. 10.
driver.do_symmetrize = 1
witness.position_mean = 0. 0. -160.e-6
witness.position_std = 2.e-6 2.e-6 5.e-6
witness.injection_type = fixed_weight
witness.num_particles = 1000000
witness.total_charge = {{witness_charge}}e-9
witness.u_mean = 0. 0. 1000.
witness.u_std = 2. 2. 10.
witness.do_symmetrize = 1
plasmas.names = electron ion
electron.density(x,y,z) = 2.e22
electron.ppc = 1 1
electron.u_mean = 0.0 0.0 0.
electron.element = electron
ion.density(x,y,z) = 2.e22
ion.ppc = 1 1
ion.u_mean = 0.0 0.0 0.
ion.element = H
diagnostic.diag_type = xz
"""Defines the analysis function that runs after the simulation."""
import os
from openpmd_viewer.addons import LpaDiagnostics
import numpy as np
from scipy.constants import e
def analyze_simulation(simulation_directory, output_params):
"""Analyze the simulation output.
This method analyzes the output generated by the simulation to
obtain the value of the optimization objective and other analyzed
parameters, if specified. The value of these parameters has to be
given to the `output_params` dictionary.
Parameters
----------
simulation_directory : str
Path to the simulation folder where the output was generated.
output_params : dict
Dictionary where the value of the objectives and analyzed parameters
will be stored. There is one entry per parameter, where the key
is the name of the parameter given by the user.
Returns
-------
dict
The `output_params` dictionary with the results from the analysis.
"""
# Open simulation diagnostics.
d = LpaDiagnostics(os.path.join(simulation_directory, "diags/hdf5"))
# Get beam particles with `u_z >= 10` and transverse offset no larger than
# 15 µm in `x` and `y`.
uz, w = d.get_particle(
["uz", "w"], iteration=d.iterations[-1], species="witness"
)
# Convert charge to pC.
q = w.sum() * e * 1e12
# Analyze distribution and fill in the output data.
if len(w) < 2: # Need at least 2 particles to calculate energy spread
output_params["f"] = 0
else:
med, mad = weighted_mad(uz / 2, w)
output_params["f"] = np.sqrt(q) * med / mad / 100
output_params["charge"] = q
output_params["energy_med"] = med
output_params["energy_mad"] = mad
return output_params
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)
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