Getting Started
¶
In this notebook, we walk through the whole bbobax contract on one function, then show what an instance is by drawing several of them. Every problem in the library -- the 24 standard functions, the noisy suite, Many-Affine, Quality-Diversity compositions -- satisfies this same three-step contract:
problem = Sphere(num_dims=2) # one function at one dimension
params = problem.sample(key) # draw an instance
eval = problem.evaluate(key, x, params)
A problem is one function at one fixed dimension, because array shapes are static in JAX. An instance is one draw of everything that varies between runs of the same function -- where the optimum sits, its value, the rotations -- exactly the role instances play in COCO. Evaluation takes a key because problems may be noisy; for noiseless functions the key changes nothing.
Install¶
You will need Python 3.11 or later, and a working JAX installation. For example, you can install JAX with:
%pip install -U "jax[cuda]"
Then, install bbobax from PyPi:
%pip install -U "bbobax[notebooks]"
Import¶
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from bbobax import Rosenbrock, RosenbrockRotated, Sphere
seed = 0
key = jax.random.key(seed)
Sphere¶
The simplest of the 24, pinned down completely: optimum at the origin with value 0, no rotation. Every constructor argument below is a default worth knowing -- x_range is the $[-5, 5]^D$ box, and x_opt_range / f_opt_range bound where an instance's optimum and its value are drawn.
problem = Sphere(
num_dims=2,
x_range=(-5, 5),
x_opt_range=(0, 0),
f_opt_range=(0, 0),
sample_rotation=False,
)
Sample an instance¶
sample draws the instance: x_opt, f_opt, the rotation matrices r and q, the instance's own key, and the noise model's parameters (a noiseless problem carries Noiseless). With the ranges pinned to zero above, this instance is the textbook sphere.
key, subkey = jax.random.split(key)
params = problem.sample(subkey)
params
BBOBParams(key=Array((), dtype=key<fry>) overlaying:
[3006591476 2914868775], x_opt=Array([0., 0.], dtype=float32), f_opt=Array(0., dtype=float32), r=Array([[1., 0.],
[0., 1.]], dtype=float32), q=Array([[1., 0.],
[0., 1.]], dtype=float32), noise_model=NoiselessParams())
Sample a solution and evaluate¶
sample_x draws uniformly from the box -- a reasonable initialization for any optimizer. evaluate returns a BBOBEval; its fitness is $f(x)$, to be minimized.
key, subkey = jax.random.split(key)
x = problem.sample_x(subkey)
key, subkey = jax.random.split(key)
print(f"fitness at x: {problem.evaluate(subkey, x, params).fitness:.3f}")
print(f"fitness at x_opt: {problem.evaluate(subkey, params.x_opt, params).fitness:.3f}")
fitness at x: 18.093 fitness at x_opt: 0.000
Plot¶
One helper, reused for every function below. Evaluation is vmapped over the grid -- the same pattern an optimizer uses to evaluate a population.
def plot_problem(problem, params, x=None, log=False):
"""Contour-plot a 2-D problem instance, with its optimum marked."""
grid = jnp.linspace(*problem.x_range, 100)
X, Y = jnp.meshgrid(grid, grid)
points = jnp.stack([X.ravel(), Y.ravel()], axis=-1)
key = jax.random.key(0) # noiseless here: the key changes nothing
fitness = jax.vmap(lambda p: problem.evaluate(key, p, params).fitness)(points)
Z = fitness.reshape(X.shape)
if log:
Z = jnp.log10(Z - jnp.min(Z) + 1e-8)
plt.figure(figsize=(7, 6))
plt.contourf(X, Y, Z, levels=50, cmap="viridis")
plt.colorbar(label="log10 fitness" if log else "fitness")
plt.plot(*params.x_opt, "r*", markersize=15, label="x_opt")
if x is not None:
plt.plot(*x, "wo", markersize=8, label="x")
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")
plt.title(f"{problem.name}, one instance at D = {problem.num_dims}")
plt.legend()
plt.show()
plot_problem(problem, params, x)
Rosenbrock¶
The same three steps on a harder landscape. Rosenbrock's banana valley spans several orders of magnitude, so the plot reads best in log scale -- which is true of most BBOB functions, and is why the benchmark measures log-precision.
problem = Rosenbrock(
num_dims=2,
x_range=(-5, 5),
x_opt_range=(0, 0),
f_opt_range=(0, 0),
sample_rotation=False,
)
key, subkey = jax.random.split(key)
params = problem.sample(subkey)
key, subkey = jax.random.split(key)
x = problem.sample_x(subkey)
plot_problem(problem, params, x, log=True)
Instances¶
So far every knob was pinned so the landscape would be the textbook one. The defaults do the opposite -- they are COCO's: x_opt uniform in $[-4, 4]^D$, and a rotation drawn per instance. Sampling twice gives two instances of the same problem: same landscape structure, different placement. This is what protects benchmark results (and meta-learned optimizers) from memorizing one layout.
problem = RosenbrockRotated(num_dims=2)
for _ in range(2):
key, subkey = jax.random.split(key)
params = problem.sample(subkey)
plot_problem(problem, params, log=True)