"""Acquisition function module for active GPR optimisation.
Implements the Expected Improvement (EI) acquisition function
for selecting the next input point to evaluate.
References
----------
Jones, D. R., Schonlau, M., & Welch, W. J. (1998).
Efficient Global Optimization of Expensive Black-Box Functions.
Journal of Global Optimization, 13(4), 455-492.
https://doi.org/10.1023/A:1008306431147
"""
from collections.abc import Sequence
import torch
from torch.distributions import Normal
from torch.quasirandom import SobolEngine
from actgpr._points import parse_search_bounds
from actgpr.surrogate import GPyTorchSurrogate
# Seed for the candidate sampler used with more than one input dimension.
# Fixed so that a multi-dimensional run is reproducible out of the box, the
# same reasoning as ObjectiveFn's DEFAULT_JITTER_SEED.
DEFAULT_CANDIDATE_SEED = 25
# Candidates scored per stage when the caller does not choose. One input
# gets 500, which is already 500 per axis. More inputs get 4000: candidates
# spread over d axes, so the same count thins out quickly, and measurement
# on a three-input problem over twelve seeds showed both a better typical
# result (median best_y 0.0037 against 0.0095 at 500) and a far better
# worst one (0.024 against 0.215), for a few seconds per run.
DEFAULT_CANDIDATES_ONE_INPUT = 500
DEFAULT_CANDIDATES_SEVERAL_INPUTS = 4000
[docs]
class Acquisition:
"""Expected Improvement acquisition function for active GPR optimisation.
Stores a reference to the surrogate, search bounds, and candidate count.
Scores candidate input points and selects the next input point to evaluate.
Candidates are placed in stages: a coarse set covering the whole search
bounds, then one or two refinement sets in a shrinking box around the
best point found so far (see ``refinement_stages``).
With one input dimension both sets are evenly spaced grids. With more,
a grid is unaffordable (500 per axis is 125 million points in 3D), so
both sets are drawn from a scrambled Sobol sequence, which covers the
space evenly without growing with the number of dimensions.
Public Methods
--------------
expected_improvement(f_mean, f_var, current_best)
Compute EI scores for an array of candidate points.
find_next_input_point(current_best)
Generate candidates, score them, and return the best input point.
"""
def __init__(
self,
surrogate: GPyTorchSurrogate,
search_bounds: Sequence[float] | Sequence[Sequence[float]],
n_candidates: int | None = None,
seed: int = DEFAULT_CANDIDATE_SEED,
refinement_stages: int | None = None,
) -> None:
"""Initialize the Acquisition function.
Parameters
----------
surrogate : GPyTorchSurrogate
The fitted surrogate model used to predict f_mean and f_var.
search_bounds : sequence of (lo, hi) pairs, or a single (lo, hi) pair
The closed interval of each input dimension within which
candidates are generated. A single pair means one dimension.
n_candidates : int or None, optional
Number of candidate points scored in each stage. If None (the
default), 500 for a single input dimension and 4000 for more,
since the same count spread over several axes thins out quickly.
Raising it further mainly improves the worst case rather than the
typical one; prediction cost grows faster than linearly with it,
so a few thousand is a sensible ceiling.
seed : int, optional
Seed for the Sobol candidate sampler used with more than one
input dimension, by default 25. Unused in one dimension, where
candidates are an evenly spaced grid.
refinement_stages : int or None, optional
How many times the zoom refinement shrinks its box around the
best candidate. If None (the default), 1 for a single input
dimension and 2 for more.
Candidates thin out as dimensions are added, so one stage leaves
a box covering more than half of each axis in three dimensions,
which barely refines anything, while in one dimension it is
already 0.8% of the axis. A second stage in one dimension would
shrink the box to the spacing between the points of a converged
run, risking evaluating the same point twice, which is the
stagnation the refinement exists to prevent.
"""
self.surrogate = surrogate
self.search_bounds = parse_search_bounds(search_bounds)
self.n_dims = len(self.search_bounds)
self.n_candidates = (
n_candidates
if n_candidates is not None
else (
DEFAULT_CANDIDATES_ONE_INPUT
if self.n_dims == 1
else DEFAULT_CANDIDATES_SEVERAL_INPUTS
)
)
self.seed = seed
self.refinement_stages = (
refinement_stages if refinement_stages is not None else min(self.n_dims, 2)
)
self._lower = torch.tensor(
[lo for lo, _ in self.search_bounds], dtype=torch.float64
)
self._upper = torch.tensor(
[hi for _, hi in self.search_bounds], dtype=torch.float64
)
self._sampler = SobolEngine(self.n_dims, scramble=True, seed=seed)
# Populated by find_next_input_point for downstream use (e.g. plotting)
self.candidates: torch.Tensor | None = None
self.f_mean: torch.Tensor | None = None
self.f_var: torch.Tensor | None = None
self.ei_scores: torch.Tensor | None = None
# The surrogate's posterior mean at the returned next_point itself
# (post zoom-refinement), distinct from f_mean, which covers only
# the coarse candidates and may not include next_point exactly.
self.next_point_mean: float | None = None
def _candidates(self, lower: torch.Tensor, upper: torch.Tensor) -> torch.Tensor:
"""Place n_candidates input points inside the box [lower, upper]."""
if self.n_dims == 1:
grid = torch.linspace(
lower.item(), upper.item(), self.n_candidates, dtype=torch.float64
)
return grid.unsqueeze(-1)
unit = self._sampler.draw(self.n_candidates, dtype=torch.float64)
return lower + (upper - lower) * unit
[docs]
def expected_improvement(
self,
f_mean: torch.Tensor,
f_var: torch.Tensor,
current_best: float,
) -> torch.Tensor:
"""Compute the Expected Improvement at candidate points.
Parameters
----------
f_mean : torch.Tensor of shape (m,)
Predicted posterior mean at candidate points.
f_var : torch.Tensor of shape (m,)
Predicted posterior variance at candidate points.
current_best : float
The smallest objective value observed so far.
Returns
-------
torch.Tensor of shape (m,)
The EI score for each candidate point.
Raises
------
ValueError
If f_mean and f_var have different shapes.
"""
if f_mean.shape != f_var.shape:
raise ValueError(
f"Shape mismatch: f_mean shape {f_mean.shape} must match "
f"f_var shape {f_var.shape}"
)
# Convert variance to standard deviation: σ = √(σ²)
# The EI formula (Jones et al. 1998) uses σ, not σ²
f_std = torch.sqrt(f_var)
# Initialise EI scores to zero for all candidates.
# Points where f_std == 0 (training points, full certainty) stay at
# EI = 0 because there is no uncertainty and thus no expected gain.
ei = torch.zeros_like(f_std)
# Boolean mask: only compute EI where f_std > 0 to avoid division by zero
mask = f_std > 0
# How much better each candidate's predicted mean is compared to current_best.
# Positive improvement means the model predicts this candidate is better.
improvement = current_best - f_mean[mask]
# Standard normal distribution N(0,1) for computing Φ (CDF) and φ (PDF)
normal = Normal(0, 1)
# Closed-form EI formula (Jones et al. 1998):
# EI(x) = (f_best − μ(x)) · Φ(z) + σ(x) · φ(z)
#
# where z = (f_best − μ(x)) / σ(x) (= improvement / f_std)
#
# Term 1: improvement * Φ(z) → exploitation
# Rewards points where the model confidently predicts a better value.
# Φ(z) is the probability that the true value is below current_best.
#
# Term 2: σ(x) * φ(z) → exploration
# Rewards points where the model is uncertain (large f_std).
# φ(z) is the bell curve height, computed as exp(log_prob(z)) because
# PyTorch's Normal distribution has no .pdf() method, only .log_prob().
# exp(log(φ(z))) = φ(z) = (1/√(2π)) · e^(−z²/2)
z = improvement / f_std[mask]
ei[mask] = improvement * normal.cdf(z) + f_std[mask] * torch.exp(
normal.log_prob(z)
)
# EI must be non-negative
assert torch.all(ei >= -1e-6), f"EI contains negative values: {ei[ei < -1e-6]}"
return ei
def __repr__(self) -> str:
"""Return a concise human-readable summary of the Acquisition."""
return (
f"Acquisition(method=EI, bounds={list(self.search_bounds)}, "
f"n_candidates={self.n_candidates})"
)