+#!/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"""
+A registry of helpers for generating inputs to acquisition function
+constructors programmatically from a consistent input format.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Hashable, Iterable, Sequence
+from typing import Any, TypeVar, Union
+
+import torch
+from botorch.acquisition.acquisition import AcquisitionFunction
+from botorch.acquisition.active_learning import qNegIntegratedPosteriorVariance
+from botorch.acquisition.analytic import (
+ ExpectedImprovement,
+ LogExpectedImprovement,
+ LogNoisyExpectedImprovement,
+ LogProbabilityOfImprovement,
+ NoisyExpectedImprovement,
+ PosteriorMean,
+ ProbabilityOfImprovement,
+ UpperConfidenceBound,
+)
+from botorch.acquisition.bayesian_active_learning import (
+ qBayesianActiveLearningByDisagreement,
+)
+from botorch.acquisition.cost_aware import InverseCostWeightedUtility
+from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction
+from botorch.acquisition.joint_entropy_search import qJointEntropySearch
+from botorch.acquisition.knowledge_gradient import (
+ qKnowledgeGradient,
+ qMultiFidelityKnowledgeGradient,
+)
+from botorch.acquisition.logei import (
+ qLogExpectedImprovement,
+ qLogNoisyExpectedImprovement,
+ TAU_MAX,
+ TAU_RELU,
+)
+from botorch.acquisition.max_value_entropy_search import (
+ qMaxValueEntropy,
+ qMultiFidelityMaxValueEntropy,
+)
+from botorch.acquisition.monte_carlo import (
+ qExpectedImprovement,
+ qLowerConfidenceBound,
+ qNoisyExpectedImprovement,
+ qProbabilityOfImprovement,
+ qSimpleRegret,
+ qUpperConfidenceBound,
+)
+from botorch.acquisition.multi_objective import (
+ ExpectedHypervolumeImprovement,
+ MCMultiOutputObjective,
+ qExpectedHypervolumeImprovement,
+ qNoisyExpectedHypervolumeImprovement,
+)
+from botorch.acquisition.multi_objective.hypervolume_knowledge_gradient import (
+ _get_hv_value_function,
+ qHypervolumeKnowledgeGradient,
+ qMultiFidelityHypervolumeKnowledgeGradient,
+)
+from botorch.acquisition.multi_objective.logei import (
+ qLogExpectedHypervolumeImprovement,
+ qLogNoisyExpectedHypervolumeImprovement,
+)
+from botorch.acquisition.multi_objective.objective import IdentityMCMultiOutputObjective
+from botorch.acquisition.multi_objective.parego import qLogNParEGO
+from botorch.acquisition.multi_objective.utils import get_default_partitioning_alpha
+from botorch.acquisition.objective import (
+ ConstrainedMCObjective,
+ IdentityMCObjective,
+ LearnedObjective,
+ MCAcquisitionObjective,
+ PosteriorTransform,
+ ScalarizedPosteriorTransform,
+)
+from botorch.acquisition.preference import (
+ AnalyticExpectedUtilityOfBestOption,
+ qExpectedUtilityOfBestOption,
+)
+from botorch.acquisition.risk_measures import RiskMeasureMCObjective
+from botorch.acquisition.utils import (
+ compute_best_feasible_objective,
+ expand_trace_observations,
+ get_infeasible_cost,
+ get_optimal_samples,
+ project_to_target_fidelity,
+)
+from botorch.exceptions.errors import UnsupportedError
+from botorch.models.cost import AffineFidelityCostModel
+from botorch.models.deterministic import FixedSingleSampleModel
+from botorch.models.gpytorch import GPyTorchModel
+from botorch.models.model import Model
+from botorch.optim.optimize import optimize_acqf
+from botorch.sampling.base import MCSampler
+from botorch.sampling.normal import IIDNormalSampler, SobolQMCNormalSampler
+from botorch.utils.containers import BotorchContainer
+from botorch.utils.datasets import SupervisedDataset
+from botorch.utils.multi_objective.box_decompositions.non_dominated import (
+ FastNondominatedPartitioning,
+ NondominatedPartitioning,
+)
+from botorch.utils.sampling import draw_sobol_samples
+from torch import Tensor
+
+
+ACQF_INPUT_CONSTRUCTOR_REGISTRY = {}
+
+T = TypeVar("T")
+MaybeDict = Union[T, dict[Hashable, T]]
+TOptimizeObjectiveKwargs = Union[
+ None,
+ MCAcquisitionObjective,
+ PosteriorTransform,
+ tuple[Tensor, Tensor],
+ dict[int, float],
+ bool,
+ int,
+ dict[str, Any],
+ Callable[[Tensor], Tensor],
+ Tensor,
+]
+
+
+def _field_is_shared(
+ datasets: Iterable[SupervisedDataset] | dict[Hashable, SupervisedDataset],
+ fieldname: str,
+) -> bool:
+ r"""Determines whether or not a given field is shared by all datasets."""
+ if isinstance(datasets, dict):
+ datasets = datasets.values()
+
+ base = None
+ for dataset in datasets:
+ if not hasattr(dataset, fieldname):
+ raise AttributeError(f"{type(dataset)} object has no field `{fieldname}`.")
+
+ obj = getattr(dataset, fieldname)
+ if base is None:
+ base = obj
+ elif isinstance(base, Tensor):
+ if not torch.equal(base, obj):
+ return False
+ elif base != obj: # pragma: no cover
+ return False
+
+ return True
+
+
+def _get_dataset_field(
+ dataset: MaybeDict[SupervisedDataset],
+ fieldname: str,
+ transform: Callable[[BotorchContainer], Any] | None = None,
+ join_rule: Callable[[Sequence[Any]], Any] | None = None,
+ first_only: bool = False,
+ assert_shared: bool = False,
+) -> Any:
+ r"""Convenience method for extracting a given field from one or more datasets."""
+ if isinstance(dataset, dict):
+ if assert_shared and not _field_is_shared(dataset, fieldname):
+ raise ValueError(f"Field `{fieldname}` must be shared.")
+
+ if not first_only:
+ fields = (
+ _get_dataset_field(d, fieldname, transform) for d in dataset.values()
+ )
+ return join_rule(tuple(fields)) if join_rule else tuple(fields)
+
+ dataset = next(iter(dataset.values()))
+
+ field = getattr(dataset, fieldname)
+ return transform(field) if transform else field
+
+
+
+
+
+
+
+
[docs]
+
def allow_only_specific_variable_kwargs(f: Callable[..., T]) -> Callable[..., T]:
+
"""
+
Decorator for allowing a function to accept keyword arguments that are not
+
explicitly listed in the function signature, but only specific ones.
+
+
This decorator is applied in `acqf_input_constructor` so that all constructors
+
obtained with `acqf_input_constructor` allow keyword
+
arguments such as `training_data` and `objective`, even if they do not appear
+
in the signature of `f`. Any other keyword arguments will raise an error.
+
"""
+
allowed = {
+
# `training_data` and/or `X_baseline` are needed to compute baselines
+
# for some EI-type acquisition functions.
+
"training_data",
+
"X_baseline",
+
# Objective thresholds are needed for defining hypervolumes in
+
# multi-objective optimization.
+
"objective_thresholds",
+
# Used in input constructors for some lookahead acquisition functions
+
# such as qKnowledgeGradient.
+
"bounds",
+
}
+
+
def g(*args: Any, **kwargs: Any) -> T:
+
new_kwargs = {}
+
accepted_kwargs = inspect.signature(f).parameters.keys()
+
for k, v in kwargs.items():
+
if k in accepted_kwargs:
+
new_kwargs[k] = v
+
elif k not in allowed:
+
raise TypeError(
+
f"Unexpected keyword argument `{k}` when"
+
f" constructing input arguments for {f.__name__}."
+
)
+
return f(*args, **new_kwargs)
+
+
return g
+
[docs]
+
@acqf_input_constructor(PosteriorMean)
+
def construct_inputs_posterior_mean(
+
model: Model,
+
posterior_transform: PosteriorTransform | None = None,
+
) -> dict[str, Model | PosteriorTransform | None]:
+
r"""Construct kwargs for PosteriorMean acquisition function.
+
+
Args:
+
model: The model to be used in the acquisition function.
+
posterior_transform: The posterior transform to be used in the
+
acquisition function.
+
+
Returns:
+
A dict mapping kwarg names of the constructor to values.
+
"""
+
return {"model": model, "posterior_transform": posterior_transform}
+
[docs]
+
def get_best_f_analytic(
+
training_data: MaybeDict[SupervisedDataset],
+
posterior_transform: PosteriorTransform | None = None,
+
) -> Tensor:
+
if isinstance(training_data, dict) and not _field_is_shared(
+
training_data, fieldname="X"
+
):
+
raise NotImplementedError("Currently only block designs are supported.")
+
+
Y = _get_dataset_field(
+
training_data,
+
fieldname="Y",
+
join_rule=lambda field_tensors: torch.cat(field_tensors, dim=-1),
+
)
+
+
if posterior_transform is not None:
+
return posterior_transform.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: MaybeDict[SupervisedDataset],
+
objective: MCAcquisitionObjective | None = None,
+
posterior_transform: PosteriorTransform | None = None,
+
constraints: list[Callable[[Tensor], Tensor]] | None = None,
+
model: Model | None = None,
+
) -> Tensor:
+
"""
+
Computes the maximum value of the objective over the training data.
+
+
Args:
+
training_data: Has fields Y, which is evaluated by `objective`, and X,
+
which is used as `X_baseline`. `Y` is of shape
+
`batch_shape x q x m`.
+
objective: The objective under which to evaluate the training data. If
+
omitted, uses `IdentityMCObjective`.
+
posterior_transform: An optional PosteriorTransform to apply to `Y`
+
before computing the objective.
+
constraints: For assessing feasibility.
+
model: Used by `compute_best_feasible_objective` when there are no
+
feasible observations.
+
+
Returns:
+
A Tensor of shape `batch_shape`.
+
"""
+
if isinstance(training_data, dict) and not _field_is_shared(
+
training_data, fieldname="X"
+
):
+
raise NotImplementedError("Currently only block designs are supported.")
+
+
X_baseline = _get_dataset_field(
+
training_data,
+
fieldname="X",
+
assert_shared=True,
+
first_only=True,
+
)
+
+
Y = _get_dataset_field(
+
training_data,
+
fieldname="Y",
+
join_rule=lambda field_tensors: torch.cat(field_tensors, dim=-1),
+
) # batch_shape x q x m
+
+
if posterior_transform is not None:
+
# retain the original tensor dimension since objective expects explicit
+
# output dimension.
+
Y_dim = Y.dim()
+
Y = posterior_transform.evaluate(Y)
+
if Y.dim() < Y_dim:
+
Y = Y.unsqueeze(-1)
+
if objective is None:
+
if Y.shape[-1] > 1:
+
raise UnsupportedError(
+
"Acquisition functions require an objective when "
+
"used with multi-output models (except for multi-objective"
+
"acquisition functions)."
+
)
+
objective = IdentityMCObjective()
+
# `Y` is of shape `(batch_shape) x q x m`; `MCAcquisitionObjective`s expect
+
# inputs `sample_shape x (batch_shape) x q x m`.
+
# For most objectives, `obj` will have shape `1 x (batch_shape) x q`, but
+
# with a `LearnedObjective` it can be `num_samples x (batch_shape) x q`.
+
obj = objective(Y.unsqueeze(0), X=X_baseline)
+
obj = obj.mean(dim=0) # taking mean over monte carlo samples
+
return compute_best_feasible_objective(
+
samples=Y,
+
obj=obj,
+
constraints=constraints,
+
model=model,
+
objective=objective,
+
posterior_transform=posterior_transform,
+
X_baseline=X_baseline,
+
)
+
[docs]
+
def optimize_objective(
+
model: Model,
+
bounds: Tensor,
+
q: int,
+
acq_function: AcquisitionFunction | None = None,
+
objective: MCAcquisitionObjective | None = None,
+
posterior_transform: PosteriorTransform | None = None,
+
linear_constraints: tuple[Tensor, Tensor] | None = None,
+
fixed_features: dict[int, float] | None = None,
+
qmc: bool = True,
+
mc_samples: int = 512,
+
seed_inner: int | None = None,
+
optimizer_options: dict[str, Any] | None = None,
+
post_processing_func: Callable[[Tensor], Tensor] | None = None,
+
batch_initial_conditions: Tensor | None = None,
+
sequential: bool = False,
+
) -> tuple[Tensor, Tensor]:
+
r"""Optimize an objective under the given model.
+
+
Args:
+
model: The model to be used in the objective.
+
bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`.
+
q: The cardinality of input sets on which the objective is to be evaluated.
+
objective: The objective to optimize.
+
posterior_transform: The posterior transform to be used in the
+
acquisition function.
+
linear_constraints: A tuple of (A, b). Given `k` linear constraints on a
+
`d`-dimensional space, `A` is `k x d` and `b` is `k x 1` such that
+
`A x <= b`. (Not used by single task models).
+
fixed_features: A dictionary of feature assignments `{feature_index: value}` to
+
hold fixed during generation.
+
qmc: Toggle for enabling (qmc=1) or disabling (qmc=0) use of Quasi Monte Carlo.
+
mc_samples: Integer number of samples used to estimate Monte Carlo objectives.
+
seed_inner: Integer seed used to initialize the sampler passed to MCObjective.
+
optimizer_options: Table used to lookup keyword arguments for the optimizer.
+
post_processing_func: A function that post-processes an optimization
+
result appropriately (i.e. according to `round-trip` transformations).
+
batch_initial_conditions: A Tensor of initial values for the optimizer.
+
sequential: If False, uses joint optimization, otherwise uses sequential
+
optimization.
+
+
Returns:
+
A tuple containing the best input locations and corresponding objective values.
+
"""
+
if optimizer_options is None:
+
optimizer_options = {}
+
+
if acq_function is None:
+
if objective is None:
+
acq_function = PosteriorMean(
+
model=model, posterior_transform=posterior_transform
+
)
+
else:
+
sampler_cls = SobolQMCNormalSampler if qmc else IIDNormalSampler
+
acq_function = qSimpleRegret(
+
model=model,
+
objective=objective,
+
posterior_transform=posterior_transform,
+
sampler=sampler_cls(
+
sample_shape=torch.Size([mc_samples]), seed=seed_inner
+
),
+
)
+
+
if fixed_features:
+
acq_function = FixedFeatureAcquisitionFunction(
+
acq_function=acq_function,
+
d=bounds.shape[-1],
+
columns=list(fixed_features.keys()),
+
values=list(fixed_features.values()),
+
)
+
free_feature_dims = list(range(len(bounds)) - fixed_features.keys())
+
free_feature_bounds = bounds[:, free_feature_dims] # (2, d' <= d)
+
else:
+
free_feature_bounds = bounds
+
+
if linear_constraints is None:
+
inequality_constraints = None
+
else:
+
A, b = linear_constraints
+
inequality_constraints = []
+
k, d = A.shape
+
for i in range(k):
+
indices = A[i, :].nonzero(as_tuple=False).squeeze()
+
coefficients = -A[i, indices]
+
rhs = -b[i, 0]
+
inequality_constraints.append((indices, coefficients, rhs))
+
+
return optimize_acqf(
+
acq_function=acq_function,
+
bounds=free_feature_bounds,
+
q=q,
+
num_restarts=optimizer_options.get("num_restarts", 60),
+
raw_samples=optimizer_options.get("raw_samples", 1024),
+
options={
+
"batch_limit": optimizer_options.get("batch_limit", 8),
+
"maxiter": optimizer_options.get("maxiter", 200),
+
"nonnegative": optimizer_options.get("nonnegative", False),
+
"method": optimizer_options.get("method", "L-BFGS-B"),
+
},
+
inequality_constraints=inequality_constraints,
+
fixed_features=None, # handled inside the acquisition function
+
post_processing_func=post_processing_func,
+
batch_initial_conditions=batch_initial_conditions,
+
return_best_only=True,
+
sequential=sequential,
+
)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ }){kind=logo}
+