Source code for botorch.acquisition.decoupled
#!/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"""Abstract base module for decoupled acquisition functions."""
from __future__ import annotations
import warnings
from abc import ABC
from typing import Optional
import torch
from botorch.acquisition.acquisition import AcquisitionFunction
from botorch.exceptions import BotorchWarning
from botorch.exceptions.errors import BotorchTensorDimensionError
from botorch.logging import shape_to_str
from botorch.models.model import ModelList
from torch import Tensor
[docs]
class DecoupledAcquisitionFunction(AcquisitionFunction, ABC):
"""
Abstract base class for decoupled acquisition functions.
A decoupled acquisition function where one may intend to
evaluate a design on only a subset of the outcomes.
Typically this would be handled by fantasizing, where one
would fantasize as to what the partial observation would
be if one were to evaluate a design on the subset of
outcomes (e.g. you only fantasize at those outcomes). The
`X_evaluation_mask` specifies which outcomes should be
evaluated for each design. `X_evaluation_mask` is `q x m`,
where there are q design points in the batch and m outcomes.
In the asynchronous case, where there are n' pending points,
we need to track which outcomes each pending point should be
evaluated on. In this case, we concatenate
`X_pending_evaluation_mask` with `X_evaluation_mask` to obtain
the full evaluation_mask.
This abstract class handles generating and updating an evaluation mask,
which is a boolean tensor indicating which outcomes a given design is
being evaluated on. The evaluation mask has shape `(n' + q) x m`, where
n' is the number of pending points and the q represents the new
candidates to be generated.
If `X(_pending)_evaluation_mas`k is None, it is assumed that `X(_pending)`
will be evaluated on all outcomes.
"""
def __init__(
self, model: ModelList, X_evaluation_mask: Optional[Tensor] = None, **kwargs
) -> None:
r"""Initialize.
Args:
model: A model
X_evaluation_mask: A `q x m`-dim boolean tensor
indicating which outcomes the decoupled acquisition
function should generate new candidates for.
"""
if not isinstance(model, ModelList):
raise ValueError(f"{self.__class__.__name__} requires using a ModelList.")
super().__init__(model=model, **kwargs)
self.num_outputs = model.num_outputs
self.X_evaluation_mask = X_evaluation_mask
self.X_pending_evaluation_mask = None
self.X_pending = None
@property
def X_evaluation_mask(self) -> Optional[Tensor]:
r"""Get the evaluation indices for the new candidate."""
return self._X_evaluation_mask
@X_evaluation_mask.setter
def X_evaluation_mask(self, X_evaluation_mask: Optional[Tensor] = None) -> None:
r"""Set the evaluation indices for the new candidate."""
if X_evaluation_mask is not None:
# TODO: Add batch support
if (
X_evaluation_mask.ndim != 2
or X_evaluation_mask.shape[-1] != self.num_outputs
):
raise BotorchTensorDimensionError(
"Expected X_evaluation_mask to be `q x m`, but got shape"
f" {shape_to_str(X_evaluation_mask.shape)}."
)
self._X_evaluation_mask = X_evaluation_mask
[docs]
def set_X_pending(
self,
X_pending: Optional[Tensor] = None,
X_pending_evaluation_mask: Optional[Tensor] = None,
) -> None:
r"""Informs the AF about pending design points for different outcomes.
Args:
X_pending: A `n' x d` Tensor with `n'` `d`-dim design points that have
been submitted for evaluation but have not yet been evaluated.
X_pending_evaluation_mask: A `n' x m`-dim tensor of booleans indicating
for which outputs the pending point is being evaluated on. If
`X_pending_evaluation_mask` is `None`, it is assumed that
`X_pending` will be evaluated on all outcomes.
"""
if X_pending is not None:
if X_pending.requires_grad:
warnings.warn(
"Pending points require a gradient but the acquisition function"
" will not provide a gradient to these points.",
BotorchWarning,
)
self.X_pending = X_pending.detach().clone()
if X_pending_evaluation_mask is not None:
if (
X_pending_evaluation_mask.ndim != 2
or X_pending_evaluation_mask.shape[0] != X_pending.shape[0]
or X_pending_evaluation_mask.shape[1] != self.num_outputs
):
raise BotorchTensorDimensionError(
f"Expected `X_pending_evaluation_mask` of shape "
f"`{X_pending.shape[0]} x {self.num_outputs}`, but "
f"got {shape_to_str(X_pending_evaluation_mask.shape)}."
)
self.X_pending_evaluation_mask = X_pending_evaluation_mask
elif self.X_evaluation_mask is not None:
raise ValueError(
"If `self.X_evaluation_mask` is not None, then "
"`X_pending_evaluation_mask` must be provided."
)
else:
self.X_pending = X_pending
self.X_pending_evaluation_mask = X_pending_evaluation_mask
[docs]
def construct_evaluation_mask(self, X: Tensor) -> Optional[Tensor]:
r"""Construct the boolean evaluation mask for X and X_pending
Args:
X: A `batch_shape x n x d`-dim tensor of designs.
Returns:
A `n + n' x m`-dim tensor of booleans indicating
which outputs should be evaluated.
"""
if self.X_pending_evaluation_mask is not None:
X_evaluation_mask = self.X_evaluation_mask
if X_evaluation_mask is None:
# evaluate all objectives for X
X_evaluation_mask = torch.ones(
X.shape[-2], self.num_outputs, dtype=torch.bool, device=X.device
)
elif X_evaluation_mask.shape[0] != X.shape[-2]:
raise BotorchTensorDimensionError(
"Expected the -2 dimension of X and X_evaluation_mask to match."
)
# construct mask for X
return torch.cat(
[X_evaluation_mask, self.X_pending_evaluation_mask], dim=-2
)
return self.X_evaluation_mask
