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 [X_evaluation_mask, self.X_pending_evaluation_mask], dim=-2 ) return self.X_evaluation_mask