Source code for aepsych.acquisition.lse

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from typing import Optional, Union

import torch
from aepsych.acquisition.objective import ProbitObjective
from botorch.acquisition.input_constructors import acqf_input_constructor
from botorch.acquisition.monte_carlo import (
    MCAcquisitionFunction,
    MCAcquisitionObjective,
    MCSampler,
)
from botorch.models.model import Model
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.utils.transforms import t_batch_mode_transform
from torch import Tensor


[docs]class MCLevelSetEstimation(MCAcquisitionFunction):
    def __init__(
        self,
        model: Model,
        target: Union[float, Tensor] = 0.75,
        beta: Union[float, Tensor] = 3.84,
        objective: Optional[MCAcquisitionObjective] = None,
        sampler: Optional[MCSampler] = None,
    ) -> None:
        r"""Monte-carlo level set estimation.

        Args:
            model: A fitted model.
            target: the level set (after objective transform) to be estimated
            beta: a parameter that governs explore-exploit tradeoff
            objective: An MCAcquisitionObjective representing the link function
                (e.g., logistic or probit.) applied on the samples.
                Can be implemented via GenericMCObjective.
            sampler: The sampler used for drawing MC samples.
        """
        if sampler is None:
            sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512]))
        if objective is None:
            objective = ProbitObjective()
        super().__init__(model=model, sampler=sampler, objective=None, X_pending=None)
        self.objective = objective
        self.beta = beta
        self.target = target

[docs]    def acquisition(self, obj_samples: torch.Tensor) -> torch.Tensor:
        """Evaluate the acquisition based on objective samples.

        Usually you should not call this directly unless you are
        subclassing this class and modifying how objective samples
        are generated.

        Args:
            obj_samples (torch.Tensor): Samples from the model, transformed
                by the objective. Should be samples x batch_shape.

        Returns:
            torch.Tensor: Acquisition function at the sampled values.
        """
        mean = obj_samples.mean(dim=0)
        variance = obj_samples.var(dim=0)
        # prevent numerical issues if probit makes all the values 1 or 0
        variance = torch.clamp(variance, min=1e-5)
        delta = torch.sqrt(self.beta * variance)
        return delta - torch.abs(mean - self.target)

    @t_batch_mode_transform()
    def forward(self, X: torch.Tensor) -> torch.Tensor:
        """Evaluate the acquisition function

        Args:
            X (torch.Tensor): Points at which to evaluate.

        Returns:
            torch.Tensor: Value of the acquisition functiona at these points.
        """

        post = self.model.posterior(X)
        samples = self.sampler(post)  # num_samples x batch_shape x q x d_out
        return self.acquisition(self.objective(samples, X)).squeeze(-1)


@acqf_input_constructor(MCLevelSetEstimation)
def construct_inputs_lse(
    model,
    training_data,
    objective=None,
    target=0.75,
    beta=3.84,
    sampler=None,
    **kwargs,
):

    return {
        "model": model,
        "objective": objective,
        "target": target,
        "beta": beta,
        "sampler": sampler,
    }