Acquisition Functions

Acquisition functions determine where to sample next in the parameter space.

Base Acquisition Function

class BOBE.acquisition.AcquisitionFunction(optimizer='scipy', optimizer_options={})

Bases: object

Base class for acquisition functions.

Acquisition functions guide the selection of new points to evaluate by balancing exploration and exploitation. Subclasses must implement the fun and get_next_point methods.

name

Name of the acquisition function.

Type:

str

optimizer

Optimizer to use (‘scipy’ or ‘optax’).

Type:

str

optimizer_options

Additional options for the optimizer.

Type:

dict

acq_optimize

Optimization function (optimize_scipy or optimize_optax).

Type:

callable

name: str = 'BaseAcquisitionFunction'

__init__(optimizer='scipy', optimizer_options={})

fun(x)

get_next_point(gp, acq_kwargs={}, maxiter=500, n_restarts=8, verbose=True, early_stop_patience=25, rng=None)

Optimize the acquisition function to obtain the next point to sample.

Parameters:

• gp (GP) – Gaussian process model.

• acq_kwargs (Dict[str, Any]) – Additional arguments for the acquisition function. Default is {}.

• maxiter (int) – Maximum number of optimization iterations. Default is 500.

• n_restarts (int) – Number of random restarts for optimization. Default is 8.

• verbose (bool) – Whether to print optimization progress. Default is True.

• early_stop_patience (int) – Patience for early stopping. Default is 25.

• rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,) and best_value is the acquisition function value.

Return type:

Tuple[ndarray, float]

get_next_batch(gp, n_batch=1, acq_kwargs={}, maxiter=500, n_restarts=8, verbose=True, early_stop_patience=25, rng=None)

Get the next batch of points to sample.

Return type:

Tuple[ndarray, float]

Weighted Integrated Posterior Variance (WIPV)

class BOBE.acquisition.WIPV(optimizer='scipy', optimizer_options={})

Bases: WeightedIntegratedPosteriorBase

Weighted Integrated Posterior Variance acquisition function.

WIPV focuses on reducing uncertainty in regions weighted by posterior probability. It integrates the posterior variance over MC samples drawn from the GP posterior, making it particularly effective for Bayesian evidence estimation.

The criterion is defined as:

WIPV(x) = E_{x’ ~ p(x’ | D)}[Var[f(x) | D]]

where the expectation is over MC samples x’ from the posterior.

Parameters:

• optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

• optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'WIPV'

fun(x, gp, mc_points=None, k_train_mc=None)

Expected Improvement (EI)

class BOBE.acquisition.EI(optimizer='scipy', optimizer_options={})

Bases: AcquisitionFunction

Expected Improvement acquisition function.

EI measures the expected improvement over the current best observed value. It balances exploitation (high mean) and exploration (high uncertainty).

The EI criterion is defined as:

EI(x) = E[max(f(x) - f_best - zeta, 0)]

where f_best is the best observed value and zeta is an exploration bonus.

Parameters:

• optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

• optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'EI'

__init__(optimizer='scipy', optimizer_options={})

fun(x, gp, best_y, zeta)

Compute Expected Improvement at point x.

Parameters:

• x (jnp.ndarray) – Point at which to evaluate EI, shape (ndim,).

• gp (GP) – Gaussian process model.

• best_y (float) – Best observed function value.

• zeta (float) – Exploration bonus parameter.

Returns:

Negative expected improvement (for minimization).

Return type:

float

get_next_point(gp, acq_kwargs, maxiter=250, n_restarts=20, verbose=True, early_stop_patience=25, rng=None)

Optimize the acquisition function to obtain the next point to sample.

Parameters:

• gp (GP) – Gaussian process model.

• acq_kwargs (dict, optional) – Additional arguments for the acquisition function. Default is {}.

• maxiter (int) – Maximum number of optimization iterations. Default is 500.

• n_restarts (int) – Number of random restarts for optimization. Default is 8.

• verbose (bool) – Whether to print optimization progress. Default is True.

• early_stop_patience (int) – Patience for early stopping. Default is 25.

• rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,) and best_value is the acquisition function value.

Return type:

tuple

Log Expected Improvement (LogEI)

class BOBE.acquisition.LogEI(optimizer='scipy', optimizer_options={})

Bases: EI

Log Expected Improvement acquisition function.

LogEI computes the logarithm of the Expected Improvement, providing better numerical stability compared to EI, especially when EI values are very small. Uses advanced numerical techniques for accurate computation in extreme cases.

Parameters:

• optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

• optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

References

[1] Ament, S., et al. (2023). “Unexpected Improvements to Expected Improvement

for Bayesian Optimization.” arXiv:2310.20708.

name: str = 'LogEI'

__init__(optimizer='scipy', optimizer_options={})

fun(x, gp, best_y, zeta)

Log Expected Improvement in pure JAX. Returns positive log-EI, so you can maximize directly.

Utility Functions

Acquisition-related utility functions.

class BOBE.acquisition.WeightedIntegratedPosteriorBase(optimizer='scipy', optimizer_options={})

Bases: AcquisitionFunction

Base class for Weighted Integrated Posterior acquisition functions.

This base class provides common functionality for acquisition functions that integrate over MC samples from the GP posterior, such as WIPV and WIPStd.

Parameters:

• optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

• optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

__init__(optimizer='scipy', optimizer_options={})

get_next_point(gp, acq_kwargs, maxiter=100, n_restarts=1, verbose=True, early_stop_patience=25, rng=None)

Optimize the acquisition function to obtain the next point.

This method is shared between WIPV and WIPStd as they follow the same optimization procedure but differ only in their objective function.

Parameters:

• gp (GP) – Gaussian process model.

• acq_kwargs (dict) – Additional arguments containing ‘mc_samples’ and optionally ‘mc_points_size’.

• maxiter (int) – Maximum optimization iterations. Default is 100.

• n_restarts (int) – Number of optimization restarts. Default is 1.

• verbose (bool) – Whether to print progress. Default is True.

• early_stop_patience (int) – Early stopping patience. Default is 25.

• rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,).

Return type:

tuple

class BOBE.acquisition.WIPStd(optimizer='scipy', optimizer_options={})

Bases: WeightedIntegratedPosteriorBase

Weighted Integrated Posterior Standard Deviation acquisition function.

WIPStd is similar to WIPV but uses standard deviation instead of variance, which can provide different exploration characteristics. It integrates the posterior standard deviation over MC samples from the GP posterior.

The criterion is defined as:

WIPStd(x) = E_{x’ ~ p(x’ | D)}[Std[f(x) | D]]

Parameters:

• optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

• optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'WIPStd'

fun(x, gp, mc_points=None, k_train_mc=None)

BOBE.acquisition.get_mc_samples(gp, warmup_steps=512, num_samples=1024, thinning=4, method='NUTS', num_chains=4, np_rng=None, rng_key=None)

BOBE.acquisition.get_mc_points(mc_samples, mc_points_size=128, rng=None)