BayesianOptimization – optimizing with a surrogate model#

BayesianOptimization is an interface which allows one to use Bayesian optimization algorithms for optimizing a device. Bayesian optimization explores a given range of selected input variables and searches for a setup which minimizes a given objective. The main advantage of that approach is that it allows one to find an approximate solution to the optimization problem after relatively small number of simulations.

The search is model-based. After every evaluation, the algorithm fits a cheap approximation of the objective to all previous observations. This approximation is then used to select the next simulation, which is expected to be the most informative. The computation spent on this approximation is negligible compared to a nextnano run. Hence, the method adds these cheap computations to decrease the number of expensive simulations.

BayesianOptimization is single-objective: Metric must return exactly one value. Metric returning several values is rejected when the optimizer is constructed. Multi-objective optimization problems can be tackled by using Evolution.

Creating the optimization problem#

Assume that the quantum well needs to be optimized to match a given transition energy. The script below shows how to optimize it by using Bayesian optimization:

import numpy as np

from nextnano_optimizers.io import IO
from nextnano_optimizers.metric import Metric
from nextnano_optimizers.optimizer import BayesianOptimization

io = IO(
    input_file_path="device.nnp",
    variable_names=["well_width", "barrier_width", "alloy_fraction"],
    target_output_paths=[("bias_00000", "Quantum", "energy_spectrum.dat")],
)

TARGET_ENERGY = 0.25  # eV

def distance_to_target(datafiles):
    energies = datafiles[0].variables["Energy"].value
    transition_energy = energies[1] - energies[0]

    return np.array([(transition_energy - TARGET_ENERGY) ** 2])

metric = Metric(
    input_shape=(1,),
    output_shape=(1,),
    extraction_function=distance_to_target,
)

optimizer = BayesianOptimization(
    nextnanoio=io,
    metric=metric,
    bounds=([2.0, 2.0, 0.15], [12.0, 20.0, 0.45]),
)

bounds of variable_names defines a design space. Candidates for each simulation are always selected inside this space. Bayesian optimization does not require an initial guess and explores the entire available design space to find the best solution.

Important: The bounds are positional. The order in which you define the bounds must match the order of variable_names in IO. The bounds are defined as two lists of lower and upper bounds:

variable_names = ["well_width", "barrier_width", "alloy_fraction"]

lower bounds   = [2.0,          2.0,             0.15]
upper bounds   = [12.0,         20.0,            0.45]

well_width     -> [2.0,  12.0]   nm
barrier_width  -> [2.0,  20.0]   nm
alloy_fraction -> [0.15,  0.45]

Note

BayesianOptimization needs the optional ax-platform dependency:

pip install nextnano_optimizers[bayesian]

It is not part of the core package. If ax-platform is missing, constructing the optimizer raises an error.

How the next simulation is chosen#

The surrogate model is a cheap approximation of the expensive simulation objective. It is built based on all previous evaluations. For any point of the design space, the model can predict an expected value of the objective and uncertainty of this expectation. Both quantities are guesses and improve during the progress of the optimization.

There are two competing strategies, whose trade-off makes the method work:

  • Exploitation focuses on samples from regions of the design space where the model expects good objective values. With a high probability, the optimal solution belongs to these regions.

  • Exploration focuses on samples from regions of the design space where the model uncertainty is high. Since the model’s predictions are uncertain, good objective values may be hidden inside these regions.

An acquisition function balances exploration and exploitation to select the next point to evaluate.

It is optimized over the allowed parameter space to find the next candidate for simulation. The acquisition function is calculated entirely on the surrogate, therefore it does not require additional simulations and optimization of the acquisition is cheap.

The optimization loop can be illustrated as follows:

observations so far
       |
       v
update the surrogate model
       |
       v
choose the most informative candidate with the acquisition function
       |
       v
nextnano simulation
       |
       v
objective value from Metric
       |
       v
add the new observation
       |
       v
repeat

At the beginning of optimization, the surrogate model does not have enough collected data to be informative. Therefore, the first few points are typically selected to collect initial observations without evaluating the surrogate model and acquisition functions.

Setting the simulation budget#

The budget of optimization can be defined by using trials:

result = optimizer.run(trials=30)

One trial is one candidate design, and therefore one nextnano simulation. Trials are evaluated one after another. The model needs each observation before it can choose the next point. Hence, trials is exactly the number of simulations the entire run may cost. trials=30 means a thirty-simulation budget. The default value of the budget is 20.

Clearly, the budget must be chosen appropriately. If it is too small, the surrogate model never becomes accurate enough to be informative. If the budget is noticeably larger than the problem needs, the later trials try to refine an answer that is already good, leading to excessive duration of optimization and to possible fluctuations of the answer caused by unnecessary explorations.

Bayesian optimization involves randomness both during initialization and when selecting subsequent candidates. To make runs reproducible, specify a random seed using run(seed=...).

Handling failed simulations#

Some points inside the bounds are not suitable for simulations of the device, e.g. because of absence of convergence. Such trials raise RuntimeError and end the run.

Alternatively, the failed simulation can be penalized by a deliberately bad objective without terminating the job. The penalty is set by non_execution_penalty:

result = optimizer.run(trials=20, non_execution_penalty=1.0)

Important

The penalty is the observed objective of the failed trial. It is included in the training data of the surrogate model and, hence, provides important feedback to the algorithm. The model is informed that some part of the design space yields a bad score and should be avoided. The penalty value is set artificially and is not related to any physics. This value must be high enough to make failed simulations unattractive. On the other hand, an extremely high value may distort the model in the vicinity of the failed simulations.

It is advisable to restrict the design space to the physically meaningful bounds and to exclude non-simulatable and meaningless parts from it. In such a setup, no penalty occurs during runs and the algorithm spends every trial on a real observation.

Running the optimization#

A setup of a complete run involves only above described parts:

optimizer = BayesianOptimization(
    nextnanoio=io,
    metric=metric,
    bounds=([2.0, 2.0, 0.15], [12.0, 20.0, 0.45]),
)

result = optimizer.run(trials=20)

Each trial involves one pass of the optimization loop: the model proposes a candidate, the simulation yields the result for this candidate, Metric scores it, and the model is informed about the observation before the next candidate is chosen.

During the run, one text line is output in a terminal per trial. It contains the trial index, its objective and the best objective found so far. This text output can be suppressed by run(verbose=False); the information is written to the run’s log either way.

The run’s progress is recorded by all algorithms of nextnano optimizers in the files Simulation_*.log, Simulation_*.json and Simulation_*.dat. During the Bayesian optimization, these files contain one entry per trial. The entries are used to plot progress figures, see Plotting the progress and results.

Inspecting the result#

The best design that the algorithm has found is called the champion:

print(optimizer.champion_x)  # the input values of the best design
print(optimizer.champion_f)  # the objective value of that design

champion_x includes the input values in the order of variable_names. champion_f is the corresponding value of the objective. run() returns both quantities:

result = optimizer.run(trials=20)

result["champion_x"]
result["champion_f"]

The champion is the best candidate discovered during the actual simulations. In other words, it is always the design whose objective was actually observed, rather than predicted by the surrogate model.

Plotting the progress and results#

run(live_progress=True) visualizes the trials’ progress in a live-plotted window; plot_all() plots the recorded history once the run has finished. Both are described in Plotting the progress and results.

Costs of BayesianOptimization#

The costs of the Bayesian optimization are governed by two ingredients: the real simulations and the intermediate computations of the algorithm. The latter are meant for fitting the surrogate model and searching for the next candidate. They are typically less expensive than the simulations themselves.

Note that the trials are sequential because each new candidate depends on all previous observations. Therefore, parallelization of the optimization (similar to Evolution’s parallel_limit) is not efficient.

If the simulations are not time-consuming and costs are not critical, another optimizer may fit better, see advantages of other algorithms on LocalOptimizer, Evolution.

When to use BayesianOptimization#

Bayesian optimization is the appropriate tool when:

  • the simulations are expensive and the simulation number is the bottleneck that makes other optimization approaches (Evolution and RL) practically not feasible,

  • intermediate computations between simulations are not time-consuming,

  • the optimization problem requires a single objective,

  • the design space is bounded and has moderate dimensionality, typically containing not more than 20 variables.

Bayesian optimization is not the appropriate tool if:

  • the optimization involves a trade-off between competing objectives,

  • the dimensionality of the design space is rather large, such that the surrogate model based on a few observations becomes unreliable.

Additional information#

See BayesianOptimization for the exact argument types, the validation rules and the attributes of BayesianOptimization.

See Bayesian Optimization for the method itself, and nextnano_optimizers.plotting module for the full set of plotting options.