#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
r"""
Utilities for acquisition functions.
"""
from __future__ import annotations
import math
from typing import Callable, Optional
import torch
from botorch.acquisition.objective import (
IdentityMCObjective,
MCAcquisitionObjective,
PosteriorTransform,
)
from botorch.exceptions.errors import (
BotorchTensorDimensionError,
DeprecationError,
UnsupportedError,
)
from botorch.models.fully_bayesian import MCMC_DIM
from botorch.models.model import Model
from botorch.sampling.base import MCSampler
from botorch.sampling.get_sampler import get_sampler
from botorch.sampling.pathwise import draw_matheron_paths
from botorch.utils.objective import compute_feasibility_indicator
from botorch.utils.sampling import optimize_posterior_samples
from botorch.utils.transforms import is_ensemble, normalize_indices
from torch import Tensor
[docs]
def get_acquisition_function(*args, **kwargs) -> None:
raise DeprecationError(
"`get_acquisition_function` has been moved to `botorch.acquisition.factory`."
)
[docs]
def repeat_to_match_aug_dim(target_tensor: Tensor, reference_tensor: Tensor) -> Tensor:
"""Repeat target_tensor until it has the same first dimension as reference_tensor
This works regardless of the batch shapes and q.
This is useful as we sometimes modify sample shapes such as in LearnedObjective.
Args:
target_tensor: A `sample_size x batch_shape x q x m`-dim Tensor
reference_tensor: A `(augmented_sample * sample_size) x batch_shape x q`-dim
Tensor. `augmented_sample` could be 1.
Returns:
The content of `target_tensor` potentially repeated so that its first dimension
matches that of `reference_tensor`.
The shape will be `(augmented_sample * sample_size) x batch_shape x q x m`.
Examples:
>>> import torch
>>> target_tensor = torch.arange(3).repeat(2, 1).T
>>> target_tensor
tensor([[0, 0],
[1, 1],
[2, 2]])
>>> repeat_to_match_aug_dim(target_tensor, torch.zeros(6))
tensor([[0, 0],
[1, 1],
[2, 2],
[0, 0],
[1, 1],
[2, 2]])
"""
augmented_sample_num, remainder = divmod(
reference_tensor.shape[0], target_tensor.shape[0]
)
if remainder != 0:
raise ValueError(
"The first dimension of reference_tensor must "
"be a multiple of target_tensor's."
)
# using repeat here as obj might be constructed as
# obj.reshape(-1, *samples.shape[2:]) where the first 2 dimensions are
# of shape `augmented_samples x sample_shape`.
repeat_size = (augmented_sample_num,) + (1,) * (target_tensor.ndim - 1)
return target_tensor.repeat(*repeat_size)
[docs]
def compute_best_feasible_objective(
samples: Tensor,
obj: Tensor,
constraints: Optional[list[Callable[[Tensor], Tensor]]],
model: Optional[Model] = None,
objective: Optional[MCAcquisitionObjective] = None,
posterior_transform: Optional[PosteriorTransform] = None,
X_baseline: Optional[Tensor] = None,
infeasible_obj: Optional[Tensor] = None,
) -> Tensor:
"""Computes the largest `obj` value that is feasible under the `constraints`. If
`constraints` is None, returns the best unconstrained objective value.
When no feasible observations exist and `infeasible_obj` is not `None`, returns
`infeasible_obj` (potentially reshaped). When no feasible observations exist and
`infeasible_obj` is `None`, uses `model`, `objective`, `posterior_transform`, and
`X_baseline` to infer and return an `infeasible_obj` `M` s.t. `M < min_x f(x)`.
Args:
samples: `(sample_shape) x batch_shape x q x m`-dim posterior samples.
obj: A `(sample_shape) x batch_shape x q`-dim Tensor of MC objective values.
constraints: A list of constraint callables which map posterior samples to
a scalar. The associated constraint is considered satisfied if this
scalar is less than zero.
model: A Model, only required when there are no feasible observations.
objective: An MCAcquisitionObjective, only optionally used when there are no
feasible observations.
posterior_transform: A PosteriorTransform, only optionally used when there are
no feasible observations.
X_baseline: A `batch_shape x d`-dim Tensor of baseline points, only required
when there are no feasible observations.
infeasible_obj: A Tensor to be returned when no feasible points exist.
Returns:
A `(sample_shape) x batch_shape`-dim Tensor of best feasible objectives.
"""
if constraints is None: # unconstrained case
# we don't need to differentiate through X_baseline for now, so taking
# the regular max over the n points to get best_f is fine
with torch.no_grad():
return obj.amax(dim=-1, keepdim=False)
is_feasible = compute_feasibility_indicator(
constraints=constraints, samples=samples
) # sample_shape x batch_shape x q
if is_feasible.any(dim=-1).all():
infeasible_value = -torch.inf
elif infeasible_obj is not None:
infeasible_value = infeasible_obj.item()
else:
if model is None:
raise ValueError(
"Must specify `model` when no feasible observation exists."
)
if X_baseline is None:
raise ValueError(
"Must specify `X_baseline` when no feasible observation exists."
)
infeasible_value = _estimate_objective_lower_bound(
model=model,
objective=objective,
posterior_transform=posterior_transform,
X=X_baseline,
).item()
is_feasible = repeat_to_match_aug_dim(
target_tensor=is_feasible, reference_tensor=obj
)
obj = torch.where(is_feasible, obj, infeasible_value)
with torch.no_grad():
return obj.amax(dim=-1, keepdim=False)
def _estimate_objective_lower_bound(
model: Model,
objective: Optional[MCAcquisitionObjective],
posterior_transform: Optional[PosteriorTransform],
X: Tensor,
) -> Tensor:
"""Estimates a lower bound on the objective values by evaluating the model at convex
combinations of `X`, returning the 6-sigma lower bound of the computed statistics.
Args:
model: A fitted model.
objective: An MCAcquisitionObjective with `m` outputs.
posterior_transform: A PosteriorTransform.
X: A `n x d`-dim Tensor of design points from which to draw convex combinations.
Returns:
A `m`-dimensional Tensor of lower bounds of the objectives.
"""
convex_weights = torch.rand(
32,
X.shape[-2],
dtype=X.dtype,
device=X.device,
)
weights_sum = convex_weights.sum(dim=0, keepdim=True)
convex_weights = convex_weights / weights_sum
# infeasible cost M is such that -M < min_x f(x), thus
# 0 < min_x f(x) - (-M), so we should take -M as a lower
# bound on the best feasible objective
return -get_infeasible_cost(
X=convex_weights @ X,
model=model,
objective=objective,
posterior_transform=posterior_transform,
)
[docs]
def get_infeasible_cost(
X: Tensor,
model: Model,
objective: Optional[Callable[[Tensor, Optional[Tensor]], Tensor]] = None,
posterior_transform: Optional[PosteriorTransform] = None,
) -> Tensor:
r"""Get infeasible cost for a model and objective.
For each outcome, computes an infeasible cost `M` such that
`-M < min_x f(x)` almost always, so that feasible points are preferred.
Args:
X: A `n x d` Tensor of `n` design points to use in evaluating the
minimum. These points should cover the design space well. The more
points the better the estimate, at the expense of added computation.
model: A fitted botorch model with `m` outcomes.
objective: The objective with which to evaluate the model output.
posterior_transform: A PosteriorTransform (optional).
Returns:
An `m`-dim tensor of infeasible cost values.
Example:
>>> model = SingleTaskGP(train_X, train_Y)
>>> objective = lambda Y: Y[..., -1] ** 2
>>> M = get_infeasible_cost(train_X, model, obj)
"""
if objective is None:
def objective(Y: Tensor, X: Optional[Tensor] = None):
return Y.squeeze(-1)
posterior = model.posterior(X, posterior_transform=posterior_transform)
lb = objective(posterior.mean - 6 * posterior.variance.clamp_min(0).sqrt(), X=X)
if lb.ndim < posterior.mean.ndim:
lb = lb.unsqueeze(-1)
# Take outcome-wise min. Looping in to handle batched models.
while lb.dim() > 1:
lb = lb.min(dim=-2).values
return -(lb.clamp_max(0.0))
[docs]
def prune_inferior_points(
model: Model,
X: Tensor,
objective: Optional[MCAcquisitionObjective] = None,
posterior_transform: Optional[PosteriorTransform] = None,
constraints: Optional[list[Callable[[Tensor], Tensor]]] = None,
num_samples: int = 2048,
max_frac: float = 1.0,
sampler: Optional[MCSampler] = None,
marginalize_dim: Optional[int] = None,
) -> Tensor:
r"""Prune points from an input tensor that are unlikely to be the best point.
Given a model, an objective, and an input tensor `X`, this function returns
the subset of points in `X` that have some probability of being the best
point under the objective. This function uses sampling to estimate the
probabilities, the higher the number of points `n` in `X` the higher the
number of samples `num_samples` should be to obtain accurate estimates.
Args:
model: A fitted model. Batched models are currently not supported.
X: An input tensor of shape `n x d`. Batched inputs are currently not
supported.
objective: The objective under which to evaluate the posterior.
posterior_transform: A PosteriorTransform (optional).
constraints: A list of constraint callables which map a Tensor of posterior
samples of dimension `sample_shape x batch-shape x q x m`-dim to a
`sample_shape x batch-shape x q`-dim Tensor. The associated constraints
are satisfied if `constraint(samples) < 0`.
num_samples: The number of samples used to compute empirical
probabilities of being the best point.
max_frac: The maximum fraction of points to retain. Must satisfy
`0 < max_frac <= 1`. Ensures that the number of elements in the
returned tensor does not exceed `ceil(max_frac * n)`.
sampler: If provided, will use this customized sampler instead of
automatically constructing one with `num_samples`.
marginalize_dim: A batch dimension that should be marginalized.
For example, this is useful when using a batched fully Bayesian
model.
Returns:
A `n' x d` with subset of points in `X`, where
n' = min(N_nz, ceil(max_frac * n))
with `N_nz` the number of points in `X` that have non-zero (empirical,
under `num_samples` samples) probability of being the best point.
"""
if marginalize_dim is None and is_ensemble(model):
# TODO: Properly deal with marginalizing fully Bayesian models
marginalize_dim = MCMC_DIM
if X.ndim > 2:
# TODO: support batched inputs (req. dealing with ragged tensors)
raise UnsupportedError(
"Batched inputs `X` are currently unsupported by prune_inferior_points"
)
if X.size(-2) == 0:
raise ValueError("X must have at least one point.")
if max_frac <= 0 or max_frac > 1.0:
raise ValueError(f"max_frac must take values in (0, 1], is {max_frac}")
max_points = math.ceil(max_frac * X.size(-2))
with torch.no_grad():
posterior = model.posterior(X=X, posterior_transform=posterior_transform)
if sampler is None:
sampler = get_sampler(
posterior=posterior, sample_shape=torch.Size([num_samples])
)
samples = sampler(posterior)
if objective is None:
objective = IdentityMCObjective()
obj_vals = objective(samples, X=X)
if obj_vals.ndim > 2:
if obj_vals.ndim == 3 and marginalize_dim is not None:
if marginalize_dim < 0:
# we do this again in compute_feasibility_indicator, but that will
# have no effect since marginalize_dim will be non-negative
marginalize_dim = (
1 + normalize_indices([marginalize_dim], d=obj_vals.ndim)[0]
)
obj_vals = obj_vals.mean(dim=marginalize_dim)
else:
# TODO: support batched inputs (req. dealing with ragged tensors)
raise UnsupportedError(
"Models with multiple batch dims are currently unsupported by"
" prune_inferior_points."
)
infeas = ~compute_feasibility_indicator(
constraints=constraints,
samples=samples,
marginalize_dim=marginalize_dim,
)
if infeas.any():
# set infeasible points to worse than worst objective
# across all samples
obj_vals[infeas] = obj_vals.min() - 1
is_best = torch.argmax(obj_vals, dim=-1)
idcs, counts = torch.unique(is_best, return_counts=True)
if len(idcs) > max_points:
counts, order_idcs = torch.sort(counts, descending=True)
idcs = order_idcs[:max_points]
return X[idcs]
[docs]
def project_to_target_fidelity(
X: Tensor,
target_fidelities: Optional[dict[int, float]] = None,
d: Optional[int] = None,
) -> Tensor:
r"""Project `X` onto the target set of fidelities.
This function assumes that the set of feasible fidelities is a box, so
projecting here just means setting each fidelity parameter to its target
value. If X does not contain the fidelity dimensions, this will insert
them and set them to their target values.
Args:
X: A `batch_shape x q x (d or d-d_f)`-dim Tensor of with `q` `d` or
`d-d_f`-dim design points for each t-batch, where d_f is the
number of fidelity dimensions. If the argument `d` is not provided,
`X` must include the fidelity dimensions and have a trailing`X` must
include the fidelity dimensions and have a trailing
target_fidelities: A dictionary mapping a subset of columns of `X` (the
fidelity parameters) to their respective target fidelity value. If
omitted, assumes that the last column of X is the fidelity parameter
with a target value of 1.0.
d: The total dimension `d`.
Return:
A `batch_shape x q x d`-dim Tensor `X_proj` with fidelity parameters
projected to the provided fidelity values.
"""
if target_fidelities is None:
target_fidelities = {-1: 1.0}
if d is None:
# assume X contains the fidelity dimensions
d = X.shape[-1]
# normalize to positive indices
tfs = {k if k >= 0 else d + k: v for k, v in target_fidelities.items()}
ones = torch.ones(*X.shape[:-1], device=X.device, dtype=X.dtype)
if X.shape[-1] == d:
# X contains fidelity dimensions
# here we're looping through the feature dimension of X - this could be
# slow for large `d`, we should optimize this for that case
X_proj = torch.stack(
[X[..., i] if i not in tfs else tfs[i] * ones for i in range(d)], dim=-1
)
elif X.shape[-1] == d - len(target_fidelities):
# need to insert fidelity dimensions
cols = []
X_idx = 0
for i in range(d):
if i not in tfs:
cols.append(X[..., X_idx])
X_idx += 1
else:
cols.append(tfs[i] * ones)
X_proj = torch.stack(cols, dim=-1)
else:
raise BotorchTensorDimensionError(
"X must have a last dimension with size `d` or `d-d_f`,"
f" but got {X.shape[-1]}."
)
return X_proj
[docs]
def expand_trace_observations(
X: Tensor, fidelity_dims: Optional[list[int]] = None, num_trace_obs: int = 0
) -> Tensor:
r"""Expand `X` with trace observations.
Expand a tensor of inputs with "trace observations" that are obtained during
the evaluation of the candidate set. This is used in multi-fidelity
optimization. It can be though of as augmenting the `q`-batch with additional
points that are the expected trace observations.
Let `f_i` be the `i`-th fidelity parameter. Then this functions assumes that
for each element of the q-batch, besides the fidelity `f_i`, we will observe
additonal fidelities `f_i1, ..., f_iK`, where `K = num_trace_obs`, during
evaluation of the candidate set `X`. Specifically, this function assumes
that `f_ij = (K-j) / (num_trace_obs + 1) * f_i` for all `i`. That is, the
expansion is performed in parallel for all fidelities (it does not expand
out all possible combinations).
Args:
X: A `batch_shape x q x d`-dim Tensor of with `q` `d`-dim design points
(incl. the fidelity parameters) for each t-batch.
fidelity_dims: The indices of the fidelity parameters. If omitted,
assumes that the last column of X contains the fidelity parameters.
num_trace_obs: The number of trace observations to use.
Return:
A `batch_shape x (q + num_trace_obs x q) x d` Tensor `X_expanded` that
expands `X` with trace observations.
"""
if num_trace_obs == 0: # No need to expand if we don't use trace observations
return X
if fidelity_dims is None:
fidelity_dims = [-1]
# The general strategy in the following is to expand `X` to the desired
# shape, and then multiply it (point-wise) with a tensor of scaling factors
reps = [1] * (X.ndim - 2) + [1 + num_trace_obs, 1]
X_expanded = X.repeat(*reps) # batch_shape x (q + num_trace_obs x q) x d
scale_fac = torch.ones_like(X_expanded)
s_pad = 1 / (num_trace_obs + 1)
# tensor of num_trace_obs scaling factors equally space between 1-s_pad and s_pad
sf = torch.linspace(1 - s_pad, s_pad, num_trace_obs, device=X.device, dtype=X.dtype)
# repeat each element q times
q = X.size(-2)
sf = torch.repeat_interleave(sf, q) # num_trace_obs * q
# now expand this to num_trace_obs x q x num_fidelities
sf = sf.unsqueeze(-1).expand(X_expanded.size(-2) - q, len(fidelity_dims))
# change relevant entries of the scaling tensor
scale_fac[..., q:, fidelity_dims] = sf
return scale_fac * X_expanded
[docs]
def project_to_sample_points(X: Tensor, sample_points: Tensor) -> Tensor:
r"""Augment `X` with sample points at which to take weighted average.
Args:
X: A `batch_shape x 1 x d`-dim Tensor of with one d`-dim design points
for each t-batch.
sample_points: `p x d'`-dim Tensor (`d' < d`) of `d'`-dim sample points at
which to compute the expectation. The `d'`-dims refer to the trailing
columns of X.
Returns:
A `batch_shape x p x d` Tensor where the q-batch includes the `p` sample points.
"""
batch_shape = X.shape[:-2]
p, d_prime = sample_points.shape
X_new = X.repeat(*(1 for _ in batch_shape), p, 1) # batch_shape x p x d
X_new[..., -d_prime:] = sample_points
return X_new
[docs]
def get_optimal_samples(
model: Model,
bounds: Tensor,
num_optima: int,
raw_samples: int = 1024,
num_restarts: int = 20,
maximize: bool = True,
) -> tuple[Tensor, Tensor]:
"""Draws sample paths from the posterior and maximizes the samples using GD.
Args:
model (Model): The model from which samples are drawn.
bounds: (Tensor): Bounds of the search space. If the model inputs are
normalized, the bounds should be normalized as well.
num_optima (int): The number of paths to be drawn and optimized.
raw_samples (int, optional): The number of candidates randomly sample.
Defaults to 1024.
num_restarts (int, optional): The number of candidates to do gradient-based
optimization on. Defaults to 20.
maximize: Whether to maximize or minimize the samples.
Returns:
Tuple[Tensor, Tensor]: The optimal input locations and corresponding
outputs, x* and f*.
"""
paths = draw_matheron_paths(model, sample_shape=torch.Size([num_optima]))
optimal_inputs, optimal_outputs = optimize_posterior_samples(
paths,
bounds=bounds,
raw_samples=raw_samples,
num_restarts=num_restarts,
maximize=maximize,
)
return optimal_inputs, optimal_outputs
