#!/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"""
Some basic data transformation helpers.
"""
from __future__ import annotations
import warnings
from collections.abc import Callable
from functools import wraps
from typing import Any, TYPE_CHECKING
import torch
from botorch.utils.safe_math import logmeanexp
from torch import Tensor
if TYPE_CHECKING: # pragma: no cover
from botorch.acquisition import AcquisitionFunction
from botorch.models.model import Model
def standardize(Y: Tensor) -> Tensor:
r"""Standardizes (zero mean, unit variance) a tensor by dim=-2.
If the tensor is single-dimensional, simply standardizes the tensor.
If for some batch index all elements are equal (or if there is only a single
data point), this function will return 0 for that batch index.
Args:
Y: A `batch_shape x n x m`-dim tensor.
Returns:
The standardized `Y`.
Example:
>>> Y = torch.rand(4, 3)
>>> Y_standardized = standardize(Y)
"""
stddim = -1 if Y.dim() < 2 else -2
Y_std = Y.std(dim=stddim, keepdim=True)
Y_std = Y_std.where(Y_std >= 1e-9, torch.full_like(Y_std, 1.0))
return (Y - Y.mean(dim=stddim, keepdim=True)) / Y_std
def _update_constant_bounds(bounds: Tensor) -> Tensor:
r"""If the lower and upper bounds are identical for a dimension, set
the upper bound to lower bound + 1.
If any modification is needed, this will return a clone of the original
tensor to avoid in-place modification.
Args:
bounds: A `2 x d`-dim tensor of lower and upper bounds.
Returns:
A `2 x d`-dim tensor of updated lower and upper bounds.
"""
if (constant_dims := (bounds[1] == bounds[0])).any():
bounds = bounds.clone()
bounds[1, constant_dims] = bounds[0, constant_dims] + 1
return bounds
def normalize(X: Tensor, bounds: Tensor, update_constant_bounds: bool = True) -> Tensor:
r"""Min-max normalize X w.r.t. the provided bounds.
Args:
X: `... x d` tensor of data
bounds: `2 x d` tensor of lower and upper bounds for each of the X's d
columns.
update_constant_bounds: If `True`, update the constant bounds in order to
avoid division by zero issues. When the upper and lower bounds are
identical for a dimension, that dimension will not be scaled. Such
dimensions will only be shifted as
`new_X[..., i] = X[..., i] - bounds[0, i]`.
Returns:
A `... x d`-dim tensor of normalized data, given by
`(X - bounds[0]) / (bounds[1] - bounds[0])`. If all elements of `X`
are contained within `bounds`, the normalized values will be
contained within `[0, 1]^d`.
Example:
>>> X = torch.rand(4, 3)
>>> bounds = torch.stack([torch.zeros(3), 0.5 * torch.ones(3)])
>>> X_normalized = normalize(X, bounds)
"""
bounds = (
_update_constant_bounds(bounds=bounds) if update_constant_bounds else bounds
)
return (X - bounds[0]) / (bounds[1] - bounds[0])
def unnormalize(
X: Tensor, bounds: Tensor, update_constant_bounds: bool = True
) -> Tensor:
r"""Un-normalizes X w.r.t. the provided bounds.
Args:
X: `... x d` tensor of data
bounds: `2 x d` tensor of lower and upper bounds for each of the X's d
columns.
update_constant_bounds: If `True`, update the constant bounds in order to
avoid division by zero issues. When the upper and lower bounds are
identical for a dimension, that dimension will not be scaled. Such
dimensions will only be shifted as
`new_X[..., i] = X[..., i] + bounds[0, i]`. This is the inverse of
the behavior of `normalize` when `update_constant_bounds=True`.
Returns:
A `... x d`-dim tensor of unnormalized data, given by
`X * (bounds[1] - bounds[0]) + bounds[0]`. If all elements of `X`
are contained in `[0, 1]^d`, the un-normalized values will be
contained within `bounds`.
Example:
>>> X_normalized = torch.rand(4, 3)
>>> bounds = torch.stack([torch.zeros(3), 0.5 * torch.ones(3)])
>>> X = unnormalize(X_normalized, bounds)
"""
bounds = (
_update_constant_bounds(bounds=bounds) if update_constant_bounds else bounds
)
return X * (bounds[1] - bounds[0]) + bounds[0]
def normalize_indices(indices: list[int] | None, d: int) -> list[int] | None:
r"""Normalize a list of indices to ensure that they are positive.
Args:
indices: A list of indices (may contain negative indices for indexing
"from the back").
d: The dimension of the tensor to index.
Returns:
A normalized list of indices such that each index is between `0` and
`d-1`, or None if indices is None.
"""
if indices is None:
return indices
normalized_indices = []
for i in indices:
if i < 0:
i = i + d
if i < 0 or i > d - 1:
raise ValueError(f"Index {i} out of bounds for tensor or length {d}.")
normalized_indices.append(i)
return normalized_indices
def _verify_output_shape(acqf: Any, X: Tensor, output: Tensor) -> bool:
r"""
Performs the output shape checks for `t_batch_mode_transform`. Output shape checks
help in catching the errors due to AcquisitionFunction arguments with erroneous
return shapes before these errors propagate further down the line.
This method checks that the `output` shape matches either the t-batch shape of X
or the `batch_shape` of `acqf.model`.
Args:
acqf: The AcquisitionFunction object being evaluated.
X: The `... x q x d`-dim input tensor with an explicit t-batch.
output: The return value of `acqf.method(X, ...)`.
Returns:
True if `output` has the correct shape, False otherwise.
"""
try:
X_batch_shape = X.shape[:-2]
if output.shape == X_batch_shape:
return True
if output.shape == torch.Size() and X_batch_shape == torch.Size([1]):
# X has a batch shape of [1] which gets squeezed.
return True
# Cases with model batch shape involved.
model_b_shape = acqf.model.batch_shape
if output.shape == model_b_shape:
# Simple inputs with batched model.
return True
model_b_dim = len(model_b_shape)
if output.shape == X_batch_shape[:-model_b_dim] + model_b_shape and all(
xs in [1, ms] for xs, ms in zip(X_batch_shape[-model_b_dim:], model_b_shape)
):
# X has additional batch dimensions beyond the model batch shape.
# For a batched model, some of the input dimensions might get broadcasted
# to the model batch shape. In that case the acquisition function output
# should replace the right-most batch dim of X with the model's batch shape.
return True
return False
except (AttributeError, NotImplementedError):
# acqf does not have model or acqf.model does not define `batch_shape`
warnings.warn(
"Output shape checks failed! Expected output shape to match t-batch shape"
f"of X, but got output with shape {output.shape} for X with shape "
f"{X.shape}. Make sure that this is the intended behavior!",
RuntimeWarning,
stacklevel=3,
)
return True
def is_fully_bayesian(model: Model) -> bool:
r"""Check if at least one model is a fully Bayesian model.
Args:
model: A BoTorch model (may be a `ModelList` or `ModelListGP`)
Returns:
True if at least one model is a fully Bayesian model.
"""
from botorch.models import ModelList
if isinstance(model, ModelList):
return any(is_fully_bayesian(m) for m in model.models)
return getattr(model, "_is_fully_bayesian", False)
def is_ensemble(model: Model) -> bool:
r"""Check if at least one model is an ensemble model.
Args:
model: A BoTorch model (may be a `ModelList` or `ModelListGP`)
Returns:
True if at least one model is an ensemble model.
"""
from botorch.models import ModelList
if isinstance(model, ModelList):
return any(is_ensemble(m) for m in model.models)
return getattr(model, "_is_ensemble", False)
def t_batch_mode_transform(
expected_q: int | None = None,
assert_output_shape: bool = True,
) -> Callable[
[Callable[[AcquisitionFunction, Any], Any]],
Callable[[AcquisitionFunction, Any], Any],
]:
r"""Factory for decorators enabling consistent t-batch behavior.
This method creates decorators for instance methods to transform an input tensor
`X` to t-batch mode (i.e. with at least 3 dimensions). This assumes the tensor
has a q-batch dimension. The decorator also checks the q-batch size if `expected_q`
is provided, and the output shape if `assert_output_shape` is `True`.
Args:
expected_q: The expected q-batch size of `X`. If specified, this will raise an
AssertionError if `X`'s q-batch size does not equal expected_q.
assert_output_shape: If `True`, this will raise an AssertionError if the
output shape does not match either the t-batch shape of `X`,
or the `acqf.model.batch_shape` for acquisition functions using
batched models.
Returns:
The decorated instance method.
Example:
>>> class ExampleClass:
>>> @t_batch_mode_transform(expected_q=1)
>>> def single_q_method(self, X):
>>> ...
>>>
>>> @t_batch_mode_transform()
>>> def arbitrary_q_method(self, X):
>>> ...
"""
def decorator(
method: Callable[[AcquisitionFunction, Any], Any],
) -> Callable[[AcquisitionFunction, Any], Any]:
@wraps(method)
def decorated(
acqf: AcquisitionFunction, X: Any, *args: Any, **kwargs: Any
) -> Any:
# Allow using acquisition functions for other inputs (e.g. lists of strings)
if not isinstance(X, Tensor):
return method(acqf, X, *args, **kwargs)
if X.dim() < 2:
raise ValueError(
f"{type(acqf).__name__} requires X to have at least 2 dimensions,"
f" but received X with only {X.dim()} dimensions."
)
elif expected_q is not None and X.shape[-2] != expected_q:
raise AssertionError(
f"Expected X to be `batch_shape x q={expected_q} x d`, but"
f" got X with shape {X.shape}."
)
X_original_shape = X.shape
# add t-batch dim
X = X if X.dim() > 2 else X.unsqueeze(0)
output = method(acqf, X, *args, **kwargs)
if assert_output_shape and not _verify_output_shape(
acqf=acqf,
X=X,
output=output,
):
raise AssertionError(
"Expected the output shape to match either the t-batch shape of "
"X, or the `model.batch_shape` in the case of acquisition "
"functions using batch models; but got output with shape "
f"{output.shape} for X with shape {X.shape}."
+ (
""
if X_original_shape == X.shape
else f" Note that `X.shape` was originally {X_original_shape} "
"before the `t_batch_mode_transform` decorator added a batch "
"dimension."
)
)
return output
return decorated
return decorator
def average_over_ensemble_models(
method: Callable[[AcquisitionFunction, Any], Any],
) -> Callable[[AcquisitionFunction, Any], Any]:
"""Decorator for averaging acquisition values over ensemble models. For example,
if the model is an ensemble, `is_ensemble(model) == True` like for a SAAS model, the
acquisition value is averaged over the samples in the ensemble.
NOTE: If the class has a `_log` attribute, the acquisition value is averaged
using logmeanexp instead of mean so that the log of the averaged acquisition value
is averaged in a numerically stable way.
Args:
method: The method to be decorated, usually `forward`.
Returns:
The decorated method.
Example:
>>> # Without decorator, forward returns a `batch_shape x ensemble_shape` tensor
>>> class SimpleAcquisition:
... def forward(self, X):
... samples, obj = self._get_samples_and_objectives(X)
... # shape is `sample_sample x batch_shape x ensemble_shape x q`
... sample_acqvals = self._sample_forward(obj)
... # return shape is `batch_shape x ensemble_shape`
... return sample_acqvals.mean(dim=0).max(dim=-1)
...
>>> # With decorator, forward returns a `batch_shape`-dim tensor
>>> class EnsembleAcquisition:
... @average_over_ensemble_models
... def forward(self, X):
... ... # same as above
... # return shape through decorator is `batch_shape`
... return sample_acqvals.mean(dim=0).max(dim=-1)
"""
def decorated(acqf: AcquisitionFunction, X: Any, *args: Any, **kwargs: Any) -> Any:
output = method(acqf, X, *args, **kwargs)
if hasattr(acqf, "model") and is_ensemble(acqf.model):
output = (
output.mean(dim=-1)
if not getattr(acqf, "_log", False)
else logmeanexp(output, dim=-1)
)
return output
return decorated
def concatenate_pending_points(
method: Callable[[Any, Tensor], Any],
) -> Callable[[Any, Tensor], Any]:
r"""Decorator concatenating X_pending into an acquisition function's argument.
This decorator works on the `forward` method of acquisition functions taking
a tensor `X` as the argument. If the acquisition function has an `X_pending`
attribute (that is not `None`), this is concatenated into the input `X`,
appropriately expanding the pending points to match the batch shape of `X`.
Example:
>>> class ExampleAcquisitionFunction:
>>> @concatenate_pending_points
>>> @t_batch_mode_transform()
>>> def forward(self, X):
>>> ...
"""
@wraps(method)
def decorated(cls: Any, X: Tensor, **kwargs: Any) -> Any:
if cls.X_pending is not None:
X = torch.cat([X, match_batch_shape(cls.X_pending, X)], dim=-2)
return method(cls, X, **kwargs)
return decorated
def match_batch_shape(X: Tensor, Y: Tensor) -> Tensor:
r"""Matches the batch dimension of a tensor to that of another tensor.
Args:
X: A `batch_shape_X x q x d` tensor, whose batch dimensions that
correspond to batch dimensions of `Y` are to be matched to those
(if compatible).
Y: A `batch_shape_Y x q' x d` tensor.
Returns:
A `batch_shape_Y x q x d` tensor containing the data of `X` expanded to
the batch dimensions of `Y` (if compatible). For instance, if `X` is
`b'' x b' x q x d` and `Y` is `b x q x d`, then the returned tensor is
`b'' x b x q x d`.
Example:
>>> X = torch.rand(2, 1, 5, 3)
>>> Y = torch.rand(2, 6, 4, 3)
>>> X_matched = match_batch_shape(X, Y)
>>> X_matched.shape
torch.Size([2, 6, 5, 3])
"""
return X.expand(X.shape[: -(Y.dim())] + Y.shape[:-2] + X.shape[-2:])