Source code for botorch.acquisition.prior_guided

#!/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.


"""
Prior-Guided Acquisition Functions

References

.. [Hvarfner2022]
    C. Hvarfner, D. Stoll, A. Souza, M. Lindauer, F. Hutter, L. Nardi. PiBO:
    Augmenting Acquisition Functions with User Beliefs for Bayesian Optimization.
    ICLR 2022.
"""
from __future__ import annotations

from typing import Optional

from botorch.acquisition.acquisition import AcquisitionFunction
from botorch.acquisition.monte_carlo import SampleReducingMCAcquisitionFunction
from botorch.utils.transforms import concatenate_pending_points, t_batch_mode_transform
from torch import Tensor

from torch.nn import Module


[docs]class PriorGuidedAcquisitionFunction(AcquisitionFunction):
    r"""Class for weighting acquisition functions by a prior distribution.

    Supports MC and batch acquisition functions via
    SampleReducingAcquisitionFunction.

    See [Hvarfner2022]_ for details.
    """

    def __init__(
        self,
        acq_function: AcquisitionFunction,
        prior_module: Module,
        log: bool = False,
        prior_exponent: float = 1.0,
        X_pending: Optional[Tensor] = None,
    ) -> None:
        r"""Initialize the prior-guided acquisition function.

        Args:
            acq_function: The base acquisition function.
            prior_module: A Module that computes the probability
                (or log probability) for the provided inputs.
                `prior_module.forward` should take a `batch_shape x q`-dim
                tensor of inputs and return a `batch_shape x q`-dim tensor
                of probabilities.
            log: A boolean that should be true if the acquisition function emits a
                log-transformed value and the prior module emits a log probability.
            prior_exponent: The exponent applied to the prior. This can be used
                for example  to decay the effect the prior over time as in
                [Hvarfner2022]_.
            X_pending: `n x d` Tensor with `n` `d`-dim design points that have
                been submitted for evaluation but have not yet been evaluated.
        """
        super().__init__(model=acq_function.model)
        self.acq_func = acq_function
        self.prior_module = prior_module
        self._log = log
        self._prior_exponent = prior_exponent
        self._is_sample_reducing_af = isinstance(
            acq_function, SampleReducingMCAcquisitionFunction
        )
        self.set_X_pending(X_pending=X_pending)

[docs]    @concatenate_pending_points
    @t_batch_mode_transform()
    def forward(self, X: Tensor) -> Tensor:
        r"""Compute the acquisition function weighted by the prior."""
        # batch_shape x q
        prior = self.prior_module(X)
        if self._is_sample_reducing_af:
            # sample_shape x batch_shape x q
            af_val = self.acq_func._non_reduced_forward(X)
        else:
            if prior.shape[-1] > 1:
                raise NotImplementedError(
                    "q-batches with q>1 are only supported using "
                    "SampleReducingMCAcquisitionFunction."
                )
            # batch_shape x q
            af_val = self.acq_func(X).unsqueeze(-1)
        if self._log:
            weighted_af_val = af_val + prior * self._prior_exponent
        else:
            weighted_af_val = af_val * prior.pow(self._prior_exponent)
        if self._is_sample_reducing_af:
            return self.acq_func._sample_reduction(
                self.acq_func._q_reduction(weighted_af_val)
            )
        return weighted_af_val.squeeze(-1)  # squeeze q-dim