Basic Optimization using BAX¶
In this notebook we demonstrate the use of Xopt to perform Bayesian Algorithm Execution (BAX) as a means of minimizing the output of a simple test function. BAX is a generalization of Bayesian Optimization that seeks to acquire observations that provide our model with maximal information about our property of interest. In this simple example, our property of interest is the minimum function output and its location in input-space. See https://arxiv.org/pdf/2209.04587.pdf for details.
Imports and random seeding for reproducibility¶
# Ignore all warnings
import warnings
warnings.filterwarnings("ignore")
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
import torch
from xopt import Xopt
from xopt.vocs import VOCS
from xopt.generators.bayesian.bax_generator import BaxGenerator
from xopt.evaluator import Evaluator
import numpy as np
import random
#random seeds for reproducibility
rand_seed = 2
torch.manual_seed(rand_seed)
np.random.seed(rand_seed) #only affects initial random observations through Xopt
random.seed(rand_seed)
Define the test problem¶
Here we define a simple optimization problem, where we attempt to minimize the sin function in the domian [0,2*pi]. Note that the function used to evaluate the objective function takes a dictionary as input and returns a dictionary as the output.
import math
# define variables and function objectives
vocs = VOCS(
variables={"x": [0, 2 * math.pi]},
observables=["y1"],
)
# define a test function to optimize
import numpy as np
def sin_function(input_dict):
return {"y1": np.sin(input_dict["x"])}
Prepare BAX generator for Xopt¶
Create a generator that uses the ExpectedInformationGain (InfoBAX) acquisition function to perform Bayesian Optimization. Note that we use minimization on a grid, so specifying the number of mesh points can negatively impact decision making time (especially in higher dimensional feature spaces).
from xopt.generators.bayesian.bax.algorithms import GridMinimize
#Prepare BAX algorithm and generator options
algorithm = GridMinimize(n_mesh_points=50)
#construct BAX generator
generator = BaxGenerator(vocs=vocs, algorithm=algorithm)
Create Evaluator and Xopt objects¶
Create the Evaluator (which allows Xopt to interface with our test function) and finish constructing our Xopt object.
#construct evaluator
evaluator = Evaluator(function=sin_function)
#construct Xopt optimizer
X = Xopt(evaluator=evaluator, generator=generator, vocs=vocs)
Generate and evaluate initial points¶
To begin optimization, we must generate some random initial data points. The first call
to X.step()
will generate and evaluate a number of randomly points specified by the
generator. Note that if we add data to xopt before calling X.step()
by assigning
the data to X.data
, calls to X.step()
will ignore the random generation and
proceed to generating points via Bayesian optimization.
# evaluate initial points
X.random_evaluate(3)
# inspect the gathered data
X.data
x | y1 | xopt_runtime | xopt_error | |
---|---|---|---|---|
0 | 3.543749 | -0.391403 | 0.000018 | False |
1 | 6.120286 | -0.162180 | 0.000002 | False |
2 | 2.829554 | 0.306999 | 0.000001 | False |
Do bayesian optimization steps¶
To perform optimization we simply call X.step()
in a loop. This allows us to do
intermediate tasks in between optimization steps, such as examining the model and
acquisition function at each step (as we demonstrate here).
import torch
import matplotlib.pyplot as plt
n_steps = 3
# test points for plotting
test_x = torch.linspace(*X.vocs.bounds.flatten(),50).double()
for i in range(5):
# get the Gaussian process model from the generator
model = X.generator.train_model()
# get acquisition function from generator
acq = X.generator.get_acquisition(model)
# calculate model posterior and acquisition function at each test point
# NOTE: need to add a dimension to the input tensor for evaluating the
# posterior and another for the acquisition function, see
# https://botorch.org/docs/batching for details
# NOTE: we use the `torch.no_grad()` environment to speed up computation by
# skipping calculations for backpropagation
with torch.no_grad():
posterior = model.posterior(test_x.unsqueeze(1))
acq_val = acq(test_x.reshape(-1,1,1))
# get mean function and confidence regions
mean = posterior.mean
l,u = posterior.mvn.confidence_region()
# plot model and acquisition function
fig,ax = plt.subplots(3,1,sharex="all")
fig.set_size_inches(8,6)
# plot model posterior
ax[0].plot(test_x, mean, label="Posterior mean")
ax[0].fill_between(test_x, l, u,alpha=0.25, label="Posterior confidence region")
# add data to model plot
ax[0].plot(X.data["x"],X.data["y1"],"C1o", label="Training data")
# plot true function
true_f = sin_function({"x": test_x})["y1"]
ax[0].plot(test_x, true_f,'--', label="Ground truth")
if i == 0:
ax[0].legend(ncols=2)
# plot the function samples and their optima found by BAX
test_points = X.generator.algorithm_results["test_points"]
posterior_samples = X.generator.algorithm_results["posterior_samples"]
execution_paths = X.generator.algorithm_results["execution_paths"]
label1 = 'Function Samples'
label2 = 'Sample Optima'
for i in range(X.generator.algorithm.n_samples):
samples, = ax[1].plot(test_points, posterior_samples[i], c='C0', alpha=0.3,
label=label1)
optima = ax[1].scatter(*execution_paths[i], c='r', marker='x', s=80,
label=label2,zorder=10)
label1 = None
label2 = None
# add legend
if i == 0:
ax[1].legend()
# plot acquisition function
ax[2].plot(test_x, acq_val.flatten())
ax[0].set_ylabel("y1")
ax[1].set_ylabel("y1")
ax[2].set_ylabel(r"$\alpha(x)$")
ax[2].set_xlabel("x")
# do the optimization step
X.step()
# access the collected data
X.data
x | y1 | xopt_runtime | xopt_error | |
---|---|---|---|---|
0 | 3.543749 | -0.391403 | 0.000018 | False |
1 | 6.120286 | -0.162180 | 0.000002 | False |
2 | 2.829554 | 0.306999 | 0.000001 | False |
3 | 4.103268 | -0.820151 | 0.000007 | False |
4 | 4.367973 | -0.941273 | 0.000008 | False |
5 | 4.710628 | -0.999998 | 0.000007 | False |
6 | 4.762932 | -0.998723 | 0.000008 | False |
7 | 4.741298 | -0.999582 | 0.000008 | False |