Source code for botorch.acquisition.input_constructors

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

r"""
A registry of helpers for generating inputs to acquisition function
constructors programmatically from a consistent input format.
"""

from typing import Any, Callable, Dict, Optional, Tuple, Type, Union

import torch
from botorch.acquisition.acquisition import AcquisitionFunction
from botorch.acquisition.analytic import (
    ExpectedImprovement,
    PosteriorMean,
    ProbabilityOfImprovement,
    UpperConfidenceBound,
    ConstrainedExpectedImprovement,
    NoisyExpectedImprovement,
)
from botorch.acquisition.monte_carlo import (
    qExpectedImprovement,
    qNoisyExpectedImprovement,
    qProbabilityOfImprovement,
    qSimpleRegret,
    qUpperConfidenceBound,
)
from botorch.acquisition.multi_objective import (
    ExpectedHypervolumeImprovement,
    qExpectedHypervolumeImprovement,
    qNoisyExpectedHypervolumeImprovement,
)
from botorch.acquisition.multi_objective.objective import (
    IdentityAnalyticMultiOutputObjective,
    IdentityMCMultiOutputObjective,
)
from botorch.acquisition.multi_objective.utils import get_default_partitioning_alpha
from botorch.acquisition.objective import (
    AcquisitionObjective,
    IdentityMCObjective,
    ScalarizedObjective,
)
from botorch.exceptions.errors import UnsupportedError
from botorch.models.model import Model
from botorch.sampling.samplers import IIDNormalSampler, MCSampler, SobolQMCNormalSampler
from botorch.utils.constraints import get_outcome_constraint_transforms
from botorch.utils.containers import TrainingData
from botorch.utils.multi_objective.box_decompositions.non_dominated import (
    FastNondominatedPartitioning,
    NondominatedPartitioning,
)
from torch import Tensor


ACQF_INPUT_CONSTRUCTOR_REGISTRY = {}


[docs]def get_acqf_input_constructor( acqf_cls: Type[AcquisitionFunction], ) -> Callable[..., Dict[str, Any]]: r"""Get acqusition function input constructor from registry. Args: acqf_cls: The AcquisitionFunction class (not instance) for which to retrieve the input constructor. Returns: The input constructor associated with `acqf_cls`. """ if acqf_cls not in ACQF_INPUT_CONSTRUCTOR_REGISTRY: raise RuntimeError( f"Input constructor for acquisition class `{acqf_cls.__name__}` not " "registered. Use the `@acqf_input_constructor` decorator to register " "a new method." ) return ACQF_INPUT_CONSTRUCTOR_REGISTRY[acqf_cls]
[docs]def acqf_input_constructor( *acqf_cls: Type[AcquisitionFunction], ) -> Callable[..., AcquisitionFunction]: r"""Decorator for registering acquisition function input constructors. Args: acqf_cls: The AcquisitionFunction classes (not instances) for which to register the input constructor. """ for acqf_cls_ in acqf_cls: if acqf_cls_ in ACQF_INPUT_CONSTRUCTOR_REGISTRY: raise ValueError( "Cannot register duplicate arg constructor for acquisition " f"class `{acqf_cls_.__name__}`" ) def decorator(method): for acqf_cls_ in acqf_cls: _register_acqf_input_constructor( acqf_cls=acqf_cls_, input_constructor=method ) ACQF_INPUT_CONSTRUCTOR_REGISTRY[acqf_cls_] = method return method return decorator
def _register_acqf_input_constructor( acqf_cls: Type[AcquisitionFunction], input_constructor: Callable[..., Dict[str, Any]], ) -> None: ACQF_INPUT_CONSTRUCTOR_REGISTRY[acqf_cls] = input_constructor # --------------------- Input argument constructors --------------------- #
[docs]@acqf_input_constructor(PosteriorMean) def construct_inputs_analytic_base( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for basic analytic acquisition functions. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. `best_f` is extracted from here. objective: The objective to in the acquisition function. Returns: A dict mapping kwarg names of the constructor to values. """ return {"model": model, "objective": objective}
[docs]@acqf_input_constructor(ExpectedImprovement, ProbabilityOfImprovement) def construct_inputs_best_f( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, maximize: bool = True, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for the acquisition functions requiring `best_f`. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. `best_f` is extracted from here. objective: The objective to in the acquisition function. maximize: If True, consider the problem a maximization problem. Returns: A dict mapping kwarg names of the constructor to values. """ base_inputs = construct_inputs_analytic_base( model=model, training_data=training_data, objective=objective ) best_f = kwargs.get( "best_f", get_best_f_analytic(training_data=training_data, objective=objective) ) return {**base_inputs, "best_f": best_f, "maximize": maximize}
[docs]@acqf_input_constructor(UpperConfidenceBound) def construct_inputs_ucb( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, beta: Union[float, Tensor] = 0.2, maximize: bool = True, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for `UpperConfidenceBound`. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. `best_f` is extracted from here. objective: The objective to in the acquisition function. beta: Either a scalar or a one-dim tensor with `b` elements (batch mode) representing the trade-off parameter between mean and covariance maximize: If True, consider the problem a maximization problem. Returns: A dict mapping kwarg names of the constructor to values. """ base_inputs = construct_inputs_analytic_base( model=model, training_data=training_data, objective=objective ) return {**base_inputs, "beta": beta, "maximize": maximize}
[docs]@acqf_input_constructor(ConstrainedExpectedImprovement) def construct_inputs_constrained_ei( model: Model, training_data: TrainingData, objective_index: int, constraints: Dict[int, Tuple[Optional[float], Optional[float]]], maximize: bool = True, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for `ConstrainedExpectedImprovement`. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. `best_f` is extracted from here. objective_index: The index of the objective. constraints: A dictionary of the form `{i: [lower, upper]}`, where `i` is the output index, and `lower` and `upper` are lower and upper bounds on that output (resp. interpreted as -Inf / Inf if None) maximize: If True, consider the problem a maximization problem. Returns: A dict mapping kwarg names of the constructor to values. """ # TODO: Implement best point computation from training data # best_f = # return { # "model": model, # "best_f": best_f, # "objective_index": objective_index, # "constraints": constraints, # "maximize": maximize, # } raise NotImplementedError # pragma: nocover
[docs]@acqf_input_constructor(NoisyExpectedImprovement) def construct_inputs_noisy_ei( model: Model, training_data: TrainingData, num_fantasies: int = 20, maximize: bool = True, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for `NoisyExpectedImprovement`. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. `best_f` is extracted from here. num_fantasies: The number of fantasies to generate. The higher this number the more accurate the model (at the expense of model complexity and performance). maximize: If True, consider the problem a maximization problem. Returns: A dict mapping kwarg names of the constructor to values. """ # TODO: Add prune_baseline functionality as for qNEI if not training_data.is_block_design: raise NotImplementedError("Currently only block designs are supported") return { "model": model, "X_observed": training_data.X, "num_fantasies": num_fantasies, "maximize": maximize, }
[docs]@acqf_input_constructor(qSimpleRegret) def construct_inputs_mc_base( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, X_pending: Optional[Tensor] = None, sampler: Optional[MCSampler] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for basic MC acquisition functions. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. Used e.g. to extract inputs such as `best_f` for expected improvement acquisition functions. objective: The objective to in the acquisition function. X_pending: A `batch_shape, m x d`-dim Tensor of `m` design points that have points that have been submitted for function evaluation but have not yet been evaluated. sampler: The sampler used to draw base samples. If omitted, uses the acquisition functions's default sampler. Returns: A dict mapping kwarg names of the constructor to values. """ return { "model": model, "objective": objective, "X_pending": X_pending, "sampler": sampler, }
[docs]@acqf_input_constructor(qExpectedImprovement) def construct_inputs_qEI( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, X_pending: Optional[Tensor] = None, sampler: Optional[MCSampler] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for the `qExpectedImprovement` constructor. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. Used e.g. to extract inputs such as `best_f` for expected improvement acquisition functions. objective: The objective to in the acquisition function. X_pending: A `m x d`-dim Tensor of `m` design points that have been submitted for function evaluation but have not yet been evaluated. Concatenated into X upon forward call. sampler: The sampler used to draw base samples. If omitted, uses the acquisition functions's default sampler. Returns: A dict mapping kwarg names of the constructor to values. """ base_inputs = construct_inputs_mc_base( model=model, training_data=training_data, objective=objective, sampler=sampler, X_pending=X_pending, ) # TODO: Dedup handling of this here and in the constructor (maybe via a # shared classmethod doing this) best_f = kwargs.get( "best_f", get_best_f_mc(training_data=training_data, objective=objective) ) return {**base_inputs, "best_f": best_f}
[docs]@acqf_input_constructor(qNoisyExpectedImprovement) def construct_inputs_qNEI( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, X_pending: Optional[Tensor] = None, sampler: Optional[MCSampler] = None, X_baseline: Optional[Tensor] = None, prune_baseline: bool = False, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for the `qNoisyExpectedImprovement` constructor. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. Used e.g. to extract inputs such as `best_f` for expected improvement acquisition functions. Only block- design training data currently supported. objective: The objective to in the acquisition function. X_pending: A `m x d`-dim Tensor of `m` design points that have been submitted for function evaluation but have not yet been evaluated. Concatenated into X upon forward call. sampler: The sampler used to draw base samples. If omitted, uses the acquisition functions's default sampler. X_baseline: A `batch_shape x r x d`-dim Tensor of `r` design points that have already been observed. These points are considered as the potential best design point. If omitted, use `training_data.X`. prune_baseline: If True, remove points in `X_baseline` that are highly unlikely to be the best point. This can significantly improve performance and is generally recommended. Returns: A dict mapping kwarg names of the constructor to values. """ base_inputs = construct_inputs_mc_base( model=model, training_data=training_data, objective=objective, sampler=sampler, X_pending=X_pending, ) if X_baseline is None: if not training_data.is_block_design: raise NotImplementedError("Currently only block designs are supported.") X_baseline = training_data.X return { **base_inputs, "X_baseline": X_baseline, "prune_baseline": prune_baseline, }
[docs]@acqf_input_constructor(qProbabilityOfImprovement) def construct_inputs_qPI( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, X_pending: Optional[Tensor] = None, sampler: Optional[MCSampler] = None, tau: float = 1e-3, best_f: Optional[Union[float, Tensor]] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for the `qProbabilityOfImprovement` constructor. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. Used e.g. to extract inputs such as `best_f` for expected improvement acquisition functions. objective: The objective to in the acquisition function. X_pending: A `m x d`-dim Tensor of `m` design points that have been submitted for function evaluation but have not yet been evaluated. Concatenated into X upon forward call. sampler: The sampler used to draw base samples. If omitted, uses the acquisition functions's default sampler. tau: The temperature parameter used in the sigmoid approximation of the step function. Smaller values yield more accurate approximations of the function, but result in gradients estimates with higher variance. best_f: The best objective value observed so far (assumed noiseless). Can be a `batch_shape`-shaped tensor, which in case of a batched model specifies potentially different values for each element of the batch. Returns: A dict mapping kwarg names of the constructor to values. """ base_inputs = construct_inputs_mc_base( model=model, training_data=training_data, objective=objective, sampler=sampler, X_pending=X_pending, ) # TODO: Dedup handling of this here and in the constructor (maybe via a # shared classmethod doing this) if best_f is None: best_f = get_best_f_mc(training_data=training_data, objective=objective) return { **base_inputs, "tau": tau, "best_f": best_f, }
[docs]@acqf_input_constructor(qUpperConfidenceBound) def construct_inputs_qUCB( model: Model, training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, X_pending: Optional[Tensor] = None, sampler: Optional[MCSampler] = None, beta: float = 0.2, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for the `qUpperConfidenceBound` constructor. Args: model: The model to be used in the acquisition function. training_data: A TrainingData object contraining the model's training data. Used e.g. to extract inputs such as `best_f` for expected improvement acquisition functions. objective: The objective to in the acquisition function. X_pending: A `m x d`-dim Tensor of `m` design points that have been submitted for function evaluation but have not yet been evaluated. Concatenated into X upon forward call. sampler: The sampler used to draw base samples. If omitted, uses the acquisition functions's default sampler. beta: Controls tradeoff between mean and standard deviation in UCB. Returns: A dict mapping kwarg names of the constructor to values. """ base_inputs = construct_inputs_mc_base( model=model, training_data=training_data, objective=objective, sampler=sampler, X_pending=X_pending, ) return {**base_inputs, "beta": beta}
def _get_sampler(mc_samples: int, qmc: bool) -> MCSampler: """Set up MC sampler for q(N)EHVI.""" # initialize the sampler seed = int(torch.randint(1, 10000, (1,)).item()) if qmc: return SobolQMCNormalSampler(num_samples=mc_samples, seed=seed) return IIDNormalSampler(num_samples=mc_samples, seed=seed)
[docs]@acqf_input_constructor(ExpectedHypervolumeImprovement) def construct_inputs_EHVI( model: Model, training_data: TrainingData, objective_thresholds: Tensor, objective: Optional[AcquisitionObjective] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for `ExpectedHypervolumeImprovement` constructor.""" num_objectives = objective_thresholds.shape[0] if kwargs.get("outcome_constraints") is not None: raise NotImplementedError("EHVI does not yet support outcome constraints.") X_observed = training_data.X alpha = kwargs.get( "alpha", get_default_partitioning_alpha(num_objectives=num_objectives), ) # This selects the objectives (a subset of the outcomes) and set each # objective threhsold to have the proper optimization direction. if objective is None: objective = IdentityAnalyticMultiOutputObjective() ref_point = objective(objective_thresholds) # Compute posterior mean (for ref point computation ref pareto frontier) # if one is not provided among arguments. Y_pmean = kwargs.get("Y_pmean") if Y_pmean is None: with torch.no_grad(): Y_pmean = model.posterior(X_observed).mean if alpha > 0: partitioning = NondominatedPartitioning( ref_point=ref_point, Y=objective(Y_pmean), alpha=alpha, ) else: partitioning = FastNondominatedPartitioning( ref_point=ref_point, Y=objective(Y_pmean), ) return { "model": model, "ref_point": ref_point, "partitioning": partitioning, "objective": objective, }
[docs]@acqf_input_constructor(qExpectedHypervolumeImprovement) def construct_inputs_qEHVI( model: Model, training_data: TrainingData, objective_thresholds: Tensor, objective: Optional[AcquisitionObjective] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for `qExpectedHypervolumeImprovement` constructor.""" X_observed = training_data.X # compute posterior mean (for ref point computation ref pareto frontier) with torch.no_grad(): Y_pmean = model.posterior(X_observed).mean outcome_constraints = kwargs.pop("outcome_constraints", None) # For HV-based acquisition functions we pass the constraint transform directly if outcome_constraints is None: cons_tfs = None else: cons_tfs = get_outcome_constraint_transforms(outcome_constraints) # Adjust `Y_pmean` to contrain feasible points only. feas = torch.stack([c(Y_pmean) <= 0 for c in cons_tfs], dim=-1).all(dim=-1) Y_pmean = Y_pmean[feas] if objective is None: objective = IdentityMCMultiOutputObjective() ehvi_kwargs = construct_inputs_EHVI( model=model, training_data=training_data, objective_thresholds=objective_thresholds, objective=objective, # Pass `Y_pmean` that accounts for constraints to `construct_inputs_EHVI` # to ensure that correct non-dominated partitioning is produced. Y_pmean=Y_pmean, **kwargs, ) sampler = kwargs.get("sampler") if sampler is None: sampler = _get_sampler( mc_samples=kwargs.get("mc_samples", 128), qmc=kwargs.get("qmc", True) ) add_qehvi_kwargs = { "sampler": sampler, "X_pending": kwargs.get("X_pending"), "constraints": cons_tfs, "eta": kwargs.get("eta", 1e-3), } return {**ehvi_kwargs, **add_qehvi_kwargs}
[docs]@acqf_input_constructor(qNoisyExpectedHypervolumeImprovement) def construct_inputs_qNEHVI( model: Model, training_data: TrainingData, objective_thresholds: Tensor, objective: Optional[AcquisitionObjective] = None, **kwargs: Any, ) -> Dict[str, Any]: r"""Construct kwargs for `qNoisyExpectedHypervolumeImprovement` constructor.""" # This selects the objectives (a subset of the outcomes) and set each # objective threhsold to have the proper optimization direction. if objective is None: objective = IdentityMCMultiOutputObjective() outcome_constraints = kwargs.pop("outcome_constraints", None) if outcome_constraints is None: cons_tfs = None else: cons_tfs = get_outcome_constraint_transforms(outcome_constraints) sampler = kwargs.get("sampler") if sampler is None: sampler = _get_sampler( mc_samples=kwargs.get("mc_samples", 128), qmc=kwargs.get("qmc", True) ) return { "model": model, "ref_point": objective(objective_thresholds), "X_baseline": kwargs.get("X_baseline", training_data.X), "sampler": sampler, "objective": objective, "constraints": cons_tfs, "X_pending": kwargs.get("X_pending"), "eta": kwargs.get("eta", 1e-3), "prune_baseline": kwargs.get("prune_baseline", True), "alpha": kwargs.get("alpha", 0.0), "cache_pending": kwargs.get("cache_pending", True), "max_iep": kwargs.get("max_iep", 0), "incremental_nehvi": kwargs.get("incremental_nehvi", True), }
[docs]def get_best_f_analytic( training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, ) -> Tensor: if not training_data.is_block_design: raise NotImplementedError("Currently only block designs are supported.") Y = training_data.Y if isinstance(objective, ScalarizedObjective): return objective.evaluate(Y).max(-1).values if Y.shape[-1] > 1: raise NotImplementedError( "Analytic acquisition functions currently only work with " "multi-output models if provided with a `ScalarizedObjective`." ) return Y.max(-2).values.squeeze(-1)
[docs]def get_best_f_mc( training_data: TrainingData, objective: Optional[AcquisitionObjective] = None, ) -> Tensor: if not training_data.is_block_design: raise NotImplementedError("Currently only block designs are supported.") Y = training_data.Y if objective is None: if Y.shape[-1] > 1: raise UnsupportedError( "Acquisition functions require an objective when " "used with multi-output models (execpt for multi-objective" "acquisition functions)." ) objective = IdentityMCObjective() return objective(training_data.Y).max(-1).values