Bayesian optimization of Branin
The problem
Branin's function, in the form Dixon and Szegö (1978) give it, is
f(x₁, x₂) = (x₂ − 5.1 x₁² / (4π²) + 5 x₁ / π − 6)² + 10 (1 − 1 / (8π)) cos x₁ + 10
on x₁ in [−5, 10] and x₂ in [0, 15]. Its three global minima, all worth 5 / (4π) ≈ 0.397887, are at (−π, 12.275), (π, 2.275) and (3π, 2.475). It's one of the problems on which Jones, Schonlau and Welch (1998) introduced efficient global optimization, the method this example runs: Bayesian optimization with the expected improvement.
Here the function stands for an expensive one, a simulation that runs for minutes or an experiment, where every evaluation counts. The question isn't how fast the search runs but how few evaluations it needs.
What makes it hard
Nothing, for a method with thousands of evaluations to spend: the Branin example finds all three minima with 30 hill climbers of 10,001 evaluations each. With tens of evaluations, every point has to be chosen with care: where the function is probably low, and where too little is known to tell. The values range from 0.4 to over 300, while the minima's basins are narrow valleys.
Representation
A Real genome of 2 genes, x₁ in [−5, 10] and x₂ in [0, 15]: genoxide's problems::Branin,
whose box and minima are those above, evaluated in Rust in both languages.
Algorithm
Bo, genoxide's Bayesian optimization, with its defaults:
- The initial design: 2(n + 1) = 6 points of a Latin hypercube, one in each sixth of each gene's range.
- The model: a Gaussian process with a constant mean and Matérn's 5/2 kernel, a length scale per gene, fitted to every evaluation so far by maximizing its marginal likelihood (Rasmussen and Williams 2006), without noise, since the function is deterministic: the model passes through every value.
- The next point: the maximum of the log expected improvement (Ament et al. 2023), the logarithm of how much a point is expected to beat the best value under the model, computed so that it keeps a gradient where the expected improvement itself underflows. It's evaluated at 1,000 random points, then maximized by L-BFGS-B from the best 10 of them and from the best point so far.
After 30 evaluations, a Gaussian process of all of them is fitted, and its posterior mean, smooth and cheap, is minimized by L-BFGS-B with its gradient from the best point found. The function is evaluated once at the result, the 31st evaluation.
Output
The first lines give the problem and the method. Then a row per evaluation: its number, the point,
the function's value and its distance above the global minimum. Generation 0 is the six points of
the design; each later row is a point the model chose. The last lines give the polish: the
point where the model's mean is lowest, the mean there, the function's value, its distance above
the nearest minimum, and the best of the 31 evaluations. In Python, run evaluates the function
in Rust and the model is genoxide's, so both versions print the same rows.
The project page plays the run back: at each step, the model's mean over the box with the points so far, the log expected improvement that chose the next point, and a curve of the best value's distance above the minimum.
Good results
A good result is within 1e-4 of a global minimum, 0.397887, in tens of evaluations. The first point the model chose, the 7th evaluation, is 2.5 above it; the 16th is 0.043 above it, near (π, 2.275), and from there the search visits all three basins, the 24th within 1.3e-3 of (−π, 12.275), the 29th within 9.6e-5 of (3π, 2.475). The model of the 30 evaluations has its lowest mean at (9.422974, 2.475867), 0.397917, and the function there is 0.397909, 2.1e-5 above the minimum: 31 evaluations in all.
Over seeds 1 to 20, the search reaches 1e-3 in a median of 30 evaluations, every seed within 80. After 30 evaluations, 4 searches are within 1e-4, and 14 with the polish; after 35, 18, and 20 with the polish. The polish costs one evaluation: the model, which passes through every value, is most precise near the best points, where the search has evaluated most.
Known optimum: 5 / (4π) ≈ 0.397887 (at three points)
Source: examples/bayesian_optimization
Interactive run: tachsin.gr/projects/genoxide/examples/bayesian-optimization
cargo run --release --example bayesian_optimization
//! Bayesian optimization of Branin's function: 30 evaluations chosen by a Gaussian process and the
//! log expected improvement, then the model's mean minimized by L-BFGS-B and evaluated once.
//!
//! The search comes within 1e-4 of one of the three global minima, and the polish of the model,
//! for one evaluation more, closes most of what is left.
//!
//! With `GENOXIDE_TRACE=<file>`, it also writes a trace of its run for the plot on the example's
//! page, with `trace.rs`.
//!
//! ```text
//! cargo run --release --example bayesian_optimization
//! ```
mod trace;
use genoxide::model::gp::GaussianProcess;
use genoxide::prelude::*;
use genoxide::problems::{Branin, Problem};
use std::cell::RefCell;
// the evaluations of the search: the initial design, then the points the model chooses
const EVALUATIONS: u64 = 30;
fn main() -> Result<()> {
let problem = Branin;
let optimum = problem.optimum().expect("known");
let minimum = optimum.value();
let bo = Bo::builder(problem.representation())
.minimize()
.seed(1)
.build()?;
println!("Branin's function in [-5, 10] x [0, 15]: three global minima of {minimum:.6}");
println!(
"{} points of a Latin hypercube, then a point per step by log-EI on a Gaussian process",
bo.initial_points()
);
println!("evaluation x1 x2 f f - f*");
// with GENOXIDE_TRACE=<file>, a trace of the run for the plot on the example's page
let trace = RefCell::new(trace::Trace::from_env());
let mut printed = 0;
let mut engine = Engine::new(bo, problem)
.stop_when(Stop::evaluations(EVALUATIONS))
.on_generation(|snapshot| {
// the points evaluated in this generation
let population = snapshot.population().as_slice();
for (index, individual) in population.iter().enumerate().skip(printed) {
let x = individual.genome();
let value = individual
.fitness()
.and_then(Fitness::score)
.expect("valid");
println!(
"{:>10} {:>10.6} {:>10.6} {:>12.6} {:>10}",
index + 1,
x[0],
x[1],
value,
scientific(value - minimum)
);
}
printed = population.len();
})
.control(|bo, progress| {
trace.borrow_mut().record(bo, progress);
Ok(())
});
let outcome = engine.run()?;
let best = outcome.best_genome().clone();
let best_value = outcome.best_fitness().score().expect("valid");
// a Gaussian process of every evaluation, its mean minimized from the best point
let evaluated = engine.algorithm().population();
let points: Vec<Reals> = evaluated.iter().map(|x| x.genome().clone()).collect();
let values: Vec<f64> = evaluated
.iter()
.map(|x| x.fitness().and_then(Fitness::score).expect("valid"))
.collect();
// the engine borrows the trace in its control
drop(engine);
let model = GaussianProcess::builder(problem.representation()).fit(&points, &values)?;
let mean = Differentiable(|x: &Reals, gradient: &mut [f64]| {
let mut variance_gradient = [0.0; 2];
let prediction = model.predict_with_gradient(x, gradient, &mut variance_gradient);
prediction.mean()
});
let lbfgsb = Lbfgsb::builder(problem.representation())
.initial_genome(best)
.minimize()
.seed(1)
.build()?;
let polished = Engine::new(lbfgsb, mean)
.stop_when(Stop::evaluations(1_000))
.run()?;
let x = polished.best_genome();
// one evaluation of the function there
let value = problem.evaluate(x);
let nearest = optimum
.solutions()
.iter()
.min_by(|a, b| distance(a, x).total_cmp(&distance(b, x)))
.expect("three minima");
println!(
"the model of the {} evaluations, its mean minimized by L-BFGS-B from the best point:",
points.len()
);
println!(
"({:.6}, {:.6}): predicted {:.6}, evaluated {value:.6}, {} above the minimum at \
({:.6}, {:.6})",
x[0],
x[1],
model.predict(x).mean(),
scientific(value - minimum),
nearest[0],
nearest[1]
);
let best_value = best_value.min(value);
println!(
"{} evaluations: the best {} above the global minimum",
points.len() + 1,
scientific(best_value - minimum)
);
assert!(best_value - minimum <= 1e-4);
trace.into_inner().write(x, value);
Ok(())
}
fn distance(a: &[f64], b: &[f64]) -> f64 {
a.iter()
.zip(b)
.map(|(x, y)| (x - y) * (x - y))
.sum::<f64>()
.sqrt()
}
// two significant digits, e.g. 1.2e-7
fn scientific(value: f64) -> String {
format!("{value:.1e}")
}
python examples/bayesian_optimization/main.py
"""Bayesian optimization of Branin's function: 30 evaluations chosen by a Gaussian process and the
log expected improvement, then the model's mean minimized by L-BFGS-B and evaluated once.
The search comes within 1e-4 of one of the three global minima, and the polish of the model, for
one evaluation more, closes most of what is left.
With ``GENOXIDE_TRACE=<file>``, it also writes a trace of its run for the plot on the example's
page, with trace.py.
python examples/bayesian_optimization/main.py
"""
import numpy as np
import genoxide as gx
from trace import Trace
# the evaluations of the search: the initial design, then the points the model chooses
EVALUATIONS = 30
def scientific(value):
"""Two significant digits, e.g. 1.2e-7."""
mantissa, exponent = f"{value:.1e}".split("e")
return f"{mantissa}e{int(exponent)}"
problem = gx.problems.Branin()
minima = problem.optimum.solutions
minimum = problem.optimum.value
bo = gx.Bo(problem.genome, objective="minimize", seed=1)
print(f"Branin's function in [-5, 10] x [0, 15]: three global minima of {minimum:.6f}")
# the initial design's default size, 2(n + 1) for n = 2 genes
print("6 points of a Latin hypercube, then a point per step by log-EI on a Gaussian process")
print("evaluation x1 x2 f f - f*")
# with GENOXIDE_TRACE=<file>, a trace of the run for the plot on the example's page
trace = Trace(minima, minimum)
evaluated = {"points": None, "values": None}
def on_generation(progress):
# the points evaluated in this generation
printed = 0 if evaluated["points"] is None else len(evaluated["points"])
for index in range(printed, len(progress.population)):
x, value = progress.population[index], float(progress.scores[index])
print(
f"{index + 1:>10} {x[0]:>10.6f} {x[1]:>10.6f} {value:>12.6f} "
f"{scientific(value - minimum):>10}"
)
evaluated["points"], evaluated["values"] = progress.population, progress.scores
result = bo.run(problem, evaluations=EVALUATIONS, on_generation=on_generation, control=trace.record)
# a Gaussian process of every evaluation, its mean minimized from the best point
points, values = evaluated["points"], evaluated["values"]
model = gx.model.gp.GaussianProcess.fit(problem.genome, points, values)
lbfgsb = gx.Lbfgsb(problem.genome, initial_genome=result.best_genome, objective="minimize", seed=1)
polished = lbfgsb.run(
lambda x: model.predict(x)[0][0],
gradient=lambda x: model.predict_with_gradient(x)[2],
evaluations=1_000,
)
x = polished.best_genome
# one evaluation of the function there
value = float(problem(x))
nearest = minima[int(np.argmin(np.linalg.norm(minima - x, axis=1)))]
print(f"the model of the {len(points)} evaluations, its mean minimized by L-BFGS-B from the best point:")
print(
f"({x[0]:.6f}, {x[1]:.6f}): predicted {model.predict(x)[0][0]:.6f}, evaluated {value:.6f}, "
f"{scientific(value - minimum)} above the minimum at ({nearest[0]:.6f}, {nearest[1]:.6f})"
)
best = min(result.best_fitness, value)
print(f"{len(points) + 1} evaluations: the best {scientific(best - minimum)} above the global minimum")
assert best - minimum <= 1e-4
trace.write(x, value)
What it prints, from a seeded run:
Branin's function in [-5, 10] x [0, 15]: three global minima of 0.397887
6 points of a Latin hypercube, then a point per step by log-EI on a Gaussian process
evaluation x1 x2 f f - f*
1 -2.094023 11.483177 7.711008 7.3e0
2 9.566534 8.178619 31.646772 3.1e1
3 1.779636 0.917823 15.079196 1.5e1
4 3.116619 3.157308 1.145222 7.5e-1
5 -4.412273 6.407455 90.516029 9.0e1
6 5.122378 13.384503 161.386238 1.6e2
7 3.802460 1.091098 2.945049 2.5e0
8 5.754346 1.947420 18.976188 1.9e1
9 -4.083638 13.352664 6.045107 5.6e0
10 -2.253838 14.243188 19.938494 2.0e1
11 4.727562 5.297559 25.625619 2.5e1
12 10.000000 3.801580 2.580940 2.2e0
13 10.000000 0.000000 10.960889 1.1e1
14 0.890363 6.322041 18.719724 1.8e1
15 -5.000000 15.000000 17.508300 1.7e1
16 3.235854 2.197541 0.440540 4.3e-2
17 9.097758 2.736199 1.180535 7.8e-1
18 9.790067 2.374804 1.212532 8.1e-1
19 8.300793 0.000000 8.707440 8.3e0
20 2.863389 2.491949 0.767187 3.7e-1
21 -3.154554 12.141734 0.425733 2.8e-2
22 -3.156264 12.562523 0.462544 6.5e-2
23 9.492983 2.838771 0.513627 1.2e-1
24 -3.153503 12.328721 0.399197 1.3e-3
25 9.363618 2.249384 0.446294 4.8e-2
26 -3.001106 11.914866 0.493115 9.5e-2
27 3.134020 2.378650 0.407715 9.8e-3
28 3.208078 2.505644 0.498582 1.0e-1
29 9.424509 2.484571 0.397984 9.6e-5
30 -3.833119 15.000000 3.606387 3.2e0
the model of the 30 evaluations, its mean minimized by L-BFGS-B from the best point:
(9.422974, 2.475867): predicted 0.397917, evaluated 0.397909, 2.1e-5 above the minimum at (9.424778, 2.475000)
31 evaluations: the best 2.1e-5 above the global minimum