Run LibEnsemble on WarpX
LibEnsemble is a library to coordinate the concurrent evaluation of dynamic ensembles of calculations. While a WarpX simulation can provide insight in some physics, it remains a single point evaluation in the space of parameters. If you have a simulation ready for use, but would like to (i) scan over some input parameters uniformly for, e.g., a tolerance study, or (ii) have a random evaluation of the space of input parameters within a given span or (iii) tune some input parameters to optimize an output parameter, e.g., beam emittance, energy spread, etc., LibEnsemble provides these capabilities and will take care of tasks monitoring with fault tolerance on multiple platforms (LibEnsemble targets modern HPC platforms like Summit).
Scripts to run LibEnsemble on WarpX simulations can be found in
This documentation does not aim at giving a training on LibEnsemble, so please refer to the LibEnsemble documentation for technical details.
WarpX example problem for LibEnsemble study
The WarpX example is built on a 2D input file (so that 1 simulation take < 1 min) of a 2-stage laser-wakefield simulation in a boosted frame. It aims at optimizing emittance preservation in the coupling between two consecutive plasma accelerator stages. Each stage accelerates an externally-injected electron beam, which charge has been tuned to show a decent amount of Ez field flattening due to longitudinal beam loading. Each stage has a parabolic transverse profile to guide the laser pulse, and a uniform longitudinal profile with cos-shape ramps and the entrance and at the exit. A fresh laser pulse is introduced at the entrance of each stage and deleted at the exit. The beam transverse distribution is matched to the first stage and the beam is injected at the beginning of the plateau of the first stage, so that the emittance is conserved in the first stage. The two stages are separated by a few-cm gap, and a focusing lens is placed in-between. Note that this is a very low resolution simulation to show an example, so it is not close to numerical convergence.
In this example, we allow LibEnsemble to tune four input parameters:
Length of the downramp of the first stage
Longitudinal position of the focusing lens (between the two stages)
Strength of the focusing lens
Length of the downramp of the second stage
The output parameter that LibEnsemble minimizes is the beam emittance at the exit of the second stage, while making sure the charge loss is small.
The scripts provided can run on a local machine or on the Summit supercomputer at OLCF. Two options are available: random sampling of parameter space or optimization on the output parameter. For the latter, we are using the Asynchronously Parallel Optimization Solver for finding Multiple Minima APOSMM method provided by LibEnsemble.
Besides a working WarpX executable, you have to install libEnsemble and its dependencies.
You can either install all packages via conda (recommended),
conda install -c conda-forge libensemble matplotlib numpy scipy yt
or try to install the same dependencies via pip (pick one or the other; note our installation details on Summit):
libensemble matplotlib nlopt numpy pytest scipy yt
See the LibEnsemble User Guide for an overview of LibEnsemble concepts. In a nutshell, a user needs to define
A generator function
gen_fthat will generate inputs for the simulation, which can be done uniformly, randomly or using an optimizer. The generator output, i.e., the simulation input, is called
'x'. The generator is provided by LibEnsemble. When the generator is an optimizer, it takes the simulation output called
'f'as an input.
A simulation function
sim_fthat will take
'x'as an input and return a single output value
'f'. In our example,
sim_fmodifies the WarpX input file depending on
'x', launches a WarpX simulation and reads the simulation output plotfile to extract
An allocator function
alloc_fthat will feed available workers with tasks. This is provided by LibEnsemble.
The files in
This is the main LibEnsemble script. It imports
warpx_simf.py(see below), defines dictionaries for parameters of each of these objects (
gen_specsincludes lower and upper bound of each element in the input array
sim_specsrespectively) and runs LibEnsemble.
- warpx_simf.run_warpx(H, persis_info, sim_specs, libE_info)
This function runs a WarpX simulation and returns quantity 'f' as well as other physical quantities measured in the run for convenience. Status check is done periodically on the simulation, provided by LibEnsemble.
WarpX input file. Some of its parameters are modified in
(util) update one WarpX input file depending on values in
'x'for this run.
- write_sim_input.write_sim_input(input_file, x_values)
This function modifies input_file to replace the values of some parameters by values provided in x_values. Parameters ---------- input_file : WarpX input file. x_values : Specific value of input x for this run.
(util) Read WarpX plotfile and return
Return optimizing quantity 'f' and other parameters for convenience. Parameters ---------- workdir : Path to directory where the simulation ran.
(util) Read LibEnsemble output files
.pickleand plot output
'f'(and other output, just for convenience) as a function of input from all simulations.
(util) dictionaries of machine-specific parameters. For convenience, the maximum number of WarpX runs is defined there.
Submission script for LibEnsemble+WarpX on Summit. Make sure to edit this file and add your project ID for allocations.
Run the example
On Summit or for a local run, LibEnsemble can run with a Random Number Generator (easiest, good as a first try) or with an optimizer (requires python package
This is set by the variable
We hereafter assume that all Python modules required are installed and that a WarpX executable is available.
local_specs dictionary in
all_machine_specs.py to fix the path to the WarpX executable (and optionally change the number of cores and OpenMP threads), and run
python run_libensemble_on_warpx.py --comms local --nworkers 3
This is adapted to a 4-core machine, as it will use:
1 process to run LibEnsemble
1 process (among the 3 workers) to run the generator
2 processes to run 2 concurrent simulations
Run on Summit at OLCF
cp -r $HOME/warpx/Tools/LibEnsemble/* sim_directory
machine = 'summit'
all_machine_specs.pyto put the right path to the WarpX executable
falseand set the project ID