Ctrl+K

Optimization with WarpX#

Description#

This examples shows how to perform a Bayesian optimization of a laser-plasma accelerator (LPA) using WarpX simulations.

The objective function to optimize (minimize) is defined as

\[f = \epsilon_f + 100\epsilon_i \left(1.0 - \frac{Q_f}{Q_i}\right)\]

where \(\epsilon_i\) and \(\epsilon_f\) are the initial and final beam emittances, respectively, and \(Q_i\) and \(Q_f\) are the initial and final beam charges. This objective is optimized by tuning 2 parameters:

  • 'adjust_factor': parameter in the range \([0.7, 1.05]\) that scales the strength of the magnetic field between the first and second stage. The value adjust_factor=1 corresponds to a focusing strength of \(454535.7\, \mathrm{T/m}\).

  • 'lens_start': the position of the start (i.e. left edge) of the focusing plasma lens in millimetres, with range \([0.32, 0.347]\).

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 WarpX simulation uses a single GPU and 4 simulations are carried out in parallel.

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
├── analysis_script.py
└── warpx

Note that the WarpX RZ executable warpx needs to be in the optimization folder. The optimization is started by executing:

python run_example.py

The scripts needed to run this example can be seen below.

run_example.py (download)#
"""Example Bayesian optimization of a multistage LPA with Warp-X.

The Warp-X simulations are performed using the template defined in the
`template_simulation_script` file.

In addition to the objective `f`, four additional parameters
are analyzed for each simulation and included 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("adjust_factor", 0.7, 1.05)
var_2 = VaryingParameter("lens_start", 0.32, 0.347)
obj = Objective("f", minimize=True)


# Define additional parameters to analyze.
energy_std = Parameter("energy_std")
energy_avg = Parameter("energy_avg")
charge = Parameter("charge")
emittance = Parameter("emittance")


# Create generator.
gen = AxSingleFidelityGenerator(
    varying_parameters=[var_1, var_2],
    objectives=[obj],
    analyzed_parameters=[energy_std, energy_avg, charge, emittance],
    n_init=4,
)


# Create evaluator.
ev = TemplateEvaluator(
    sim_template="template_simulation_script",
    analysis_func=analyze_simulation,
    executable="warpx",
    n_gpus=1,
)


# Create exploration.
exp = Exploration(
    generator=gen, evaluator=ev, max_evals=1000, sim_workers=4, run_async=True
)


# 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()
template_simulation_script (download)#
# algo
algo.charge_deposition = standard
algo.current_deposition = direct
algo.field_gathering = momentum-conserving
algo.maxwell_solver = PSATD
algo.particle_pusher = vay
algo.particle_shape = 3

# amr
amr.blocking_factor = 64
amr.max_grid_size = 4000
amr.max_level = 0
amr.n_cell = 64 3840

# beam
beam.charge = -q_e
beam.initialize_self_fields = 0
beam.injection_style = gaussian_beam
beam.do_symmetrize = 1
beam.mass = m_e
beam.momentum_distribution_type = gaussian
beam.npart = 50000
beam.q_tot = -1.e-15
beam.rigid_advance = 1
beam.ux_m = 0.0
beam.ux_th = 1.0
beam.uy_m = 0.0
beam.uy_th = 1.0
beam.uz_m = 1956.9469069265979
beam.uz_th = 0.
beam.x_m = 0.0
beam.x_rms = 3e-06
beam.y_m = 0.0
beam.y_rms = 3e-06
beam.z_m = -0.000109
beam.z_rms = 5.e-7
beam.zinject_plane = 0.0

# boundary
boundary.field_hi = none damped
boundary.field_lo = none damped
boundary.particle_hi = absorbing absorbing
boundary.particle_lo = absorbing absorbing

# diag
diag.diag_type = Full
diag.fields_to_plot = Er Et Ez rho
diag.file_prefix = diag
diag.format = openpmd
diag.intervals = 1000
diag.species = beam

# diagnostics
diagnostics.diags_names = diag

# electrons1
electrons1.charge = -q_e
electrons1.density_function(x,y,z) = "n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.-cos(pi*(z-0.0)/Lplus)))*((z-0.0)<Lplus)+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*((z-0.0)>=Lplus)*((z-0.0)<(Lplus+Lp))+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.+cos(pi*((z-0.0)-Lplus-Lp)/Lminus)))*((z-0.0)>=(Lplus+Lp))*((z-0.0)<(Lplus+Lp+Lminus))"
electrons1.do_continuous_injection = 1
electrons1.initialize_self_fields = 0
electrons1.injection_style = nuniformpercell
electrons1.mass = m_e
electrons1.momentum_distribution_type = constant
electrons1.num_particles_per_cell_each_dim = 2 8 2
electrons1.profile = parse_density_function
electrons1.ux = 0.0
electrons1.uy = 0.0
electrons1.uz = 0.0
electrons1.xmax = 0.0001
electrons1.xmin = -0.0001
electrons1.ymax = 0.0001
electrons1.ymin = -0.0001
electrons1.zmax = 0.32
electrons1.zmin = 0.0

# electrons2
electrons2.charge = -q_e
electrons2.density_function(x,y,z) = "n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.-cos(pi*(z-0.35)/Lplus)))*((z-0.35)<Lplus)+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*((z-0.35)>=Lplus)*((z-0.35)<(Lplus+Lp))+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.+cos(pi*((z-0.35)-Lplus-Lp)/Lminus)))*((z-0.35)>=(Lplus+Lp))*((z-0.35)<(Lplus+Lp+Lminus))"
electrons2.do_continuous_injection = 1
electrons2.initialize_self_fields = 0
electrons2.injection_style = nuniformpercell
electrons2.mass = m_e
electrons2.momentum_distribution_type = constant
electrons2.num_particles_per_cell_each_dim = 2 8 2
electrons2.profile = parse_density_function
electrons2.ux = 0.0
electrons2.uy = 0.0
electrons2.uz = 0.0
electrons2.xmax = 0.0001
electrons2.xmin = -0.0001
electrons2.ymax = 0.0001
electrons2.ymin = -0.0001
electrons2.zmax = 0.6699999999999999
electrons2.zmin = 0.35

# geometry
geometry.dims = RZ
geometry.prob_hi = 128.e-6 0.
geometry.prob_lo = 0.   -180.e-6

# ions1
ions1.charge = q_e
ions1.density_function(x,y,z) = "n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.-cos(pi*(z-0.0)/Lplus)))*((z-0.0)<Lplus)+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*((z-0.0)>=Lplus)*((z-0.0)<(Lplus+Lp))+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.+cos(pi*((z-0.0)-Lplus-Lp)/Lminus)))*((z-0.0)>=(Lplus+Lp))*((z-0.0)<(Lplus+Lp+Lminus))"
ions1.do_continuous_injection = 1
ions1.initialize_self_fields = 0
ions1.injection_style = nuniformpercell
ions1.mass = m_p
ions1.momentum_distribution_type = constant
ions1.num_particles_per_cell_each_dim = 2 8 2
ions1.profile = parse_density_function
ions1.ux = 0.0
ions1.uy = 0.0
ions1.uz = 0.0
ions1.xmax = 0.0001
ions1.xmin = -0.0001
ions1.ymax = 0.0001
ions1.ymin = -0.0001
ions1.zmax = 0.32
ions1.zmin = 0.0

# ions2
ions2.charge = q_e
ions2.density_function(x,y,z) = "n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.-cos(pi*(z-0.35)/Lplus)))*((z-0.35)<Lplus)+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*((z-0.35)>=Lplus)*((z-0.35)<(Lplus+Lp))+n0*(1.+4.*(x**2+y**2)/(kp**2*Rc**4))*(0.5*(1.+cos(pi*((z-0.35)-Lplus-Lp)/Lminus)))*((z-0.35)>=(Lplus+Lp))*((z-0.35)<(Lplus+Lp+Lminus))"
ions2.do_continuous_injection = 1
ions2.initialize_self_fields = 0
ions2.injection_style = nuniformpercell
ions2.mass = m_p
ions2.momentum_distribution_type = constant
ions2.num_particles_per_cell_each_dim = 2 8 2
ions2.profile = parse_density_function
ions2.ux = 0.0
ions2.uy = 0.0
ions2.uz = 0.0
ions2.xmax = 0.0001
ions2.xmin = -0.0001
ions2.ymax = 0.0001
ions2.ymin = -0.0001
ions2.zmax = 0.6699999999999999
ions2.zmin = 0.35

# laser1
laser1.direction = 0.0 0.0 1.0
laser1.do_continuous_injection = 0
laser1.e_max = 6822740000000.0
laser1.polarization = 0.0 1.0 0.0
laser1.position = 0.0 0.0 -1e-09
laser1.profile = Gaussian
laser1.profile_duration = 7.33841e-14
laser1.profile_focal_distance = 0.00875
laser1.profile_t_peak = 1.46764864e-13
laser1.profile_waist = 5e-05
laser1.wavelength = 8e-07

# laser2
laser2.direction = 0.0 0.0 1.0
laser2.do_continuous_injection = 1
laser2.e_max = 6822740000000.0
laser2.polarization = 0.0 1.0 0.0
laser2.position = 0.0 0.0 0.34999999899999995
laser2.profile = Gaussian
laser2.profile_duration = 7.33841e-14
laser2.profile_focal_distance = 0.00874999999999998
laser2.profile_t_peak = 1.167621098057891e-09
laser2.profile_waist = 5e-05
laser2.wavelength = 8e-07

# lasers
lasers.names = laser1 laser2

# my_constants
my_constants.Lminus = 0.
my_constants.Lp = 0.32
my_constants.Lplus = 0.
my_constants.Rc = 4e-05
my_constants.kp = 77588.13070567355
my_constants.n0 = 1.7e+23
my_constants.pi = 3.141592653589793

# particles
particles.B_ext_particle_init_style = repeated_plasma_lens
particles.repeated_plasma_lens_period = 0.35
particles.repeated_plasma_lens_starts = {{lens_start}}
particles.repeated_plasma_lens_lengths = 0.003
particles.repeated_plasma_lens_strengths_B = {{adjust_factor}}*1363607.21922141/3
particles.rigid_injected_species = beam
particles.species_names = electrons1 ions1 electrons2 ions2 beam

# psatd
psatd.current_correction = 1
psatd.nox = 16
psatd.noz = 16
psatd.nx_guard = 24
psatd.nz_guard = 32
psatd.update_with_rho = 1
psatd.v_galilean = 0.0 0.0 -0.9998611014647094

# warpx
warpx.boost_direction = z
warpx.cfl = 0.9999
warpx.do_moving_window = 1
warpx.filter_npass_each_dir = 1 1
warpx.gamma_boost = 60.0
warpx.mirror_z = 0.321
warpx.mirror_z_npoints = 4
warpx.mirror_z_width = 8e-06
warpx.moving_window_dir = z
warpx.moving_window_v = 1.0
warpx.n_rz_azimuthal_modes = 2
warpx.num_mirrors = 1
warpx.use_filter = 1
warpx.use_filter_compensation = 0
warpx.verbose = 2
warpx.do_single_precision_comms = 1
warpx.zmax_plasma_to_compute_max_step = 0.7
analysis_script.py (download)#
"""Defines the analysis function that runs after the simulation."""

import os

import numpy as np
from openpmd_viewer.addons import LpaDiagnostics


def get_emittance(ts, t):
    """Calculate the beam emittance at the given time step."""
    w, x, ux = ts.get_particle(["w", "x", "ux"], t=t)
    x2 = np.average(x**2, weights=w)
    u2 = np.average(ux**2, weights=w)
    xu = np.average(x * ux, weights=w)
    return np.sqrt(x2 * u2 - xu**2)


def analyze_simulation(simulation_directory, output_params):
    """Analyze the output of the WarpX simulation.

    The function calculates the objective function 'f' as well as the
    diagnostic quantities listed as `analyzed_parameters` in the generator.
    """
    ts = LpaDiagnostics(os.path.join(simulation_directory, "diag"))
    t0 = 4.0e-11  # Time (boosted-frame) at which we calculate beam properties.

    charge_i = ts.get_charge(t=0)
    emittance_i = get_emittance(ts, t=0)
    charge_f = ts.get_charge(t=t0)
    emittance_f = get_emittance(ts, t=t0)
    energy_avg, energy_std = ts.get_mean_gamma(t=t0)

    # Here: Build a quantity to minimize (f) that encompasses
    # emittance AND charge loss 1% charge loss has the
    # same impact as doubling the initial emittance.
    # we minimize f!
    output_params["f"] = np.log(
        emittance_f + emittance_i * (1.0 - charge_f / charge_i) * 100
    )
    output_params["energy_std"] = energy_std
    output_params["energy_avg"] = energy_avg
    output_params["charge"] = charge_f
    output_params["emittance"] = emittance_f

    return output_params
On this page
Edit on GitHub
Show Source