#!/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.
from __future__ import annotations
import dataclasses
import inspect
import warnings
from logging import Logger
from typing import Any, Dict, List, Optional, Tuple, Type
import torch
from ax.core.search_space import SearchSpaceDigest
from ax.core.types import TCandidateMetadata
from ax.exceptions.core import AxWarning, UnsupportedError, UserInputError
from ax.models.model_utils import best_in_sample_point
from ax.models.torch.botorch_modular.utils import (
disable_one_to_many_transforms,
fit_botorch_model,
)
from ax.models.torch.utils import (
_to_inequality_constraints,
pick_best_out_of_sample_point_acqf_class,
predict_from_model,
)
from ax.models.torch_base import TorchOptConfig
from ax.models.types import TConfig
from ax.utils.common.base import Base
from ax.utils.common.constants import Keys
from ax.utils.common.logger import get_logger
from ax.utils.common.typeutils import checked_cast, checked_cast_optional, not_none
from botorch.models.model import Model
from botorch.models.pairwise_gp import PairwiseGP
from botorch.models.transforms.input import InputPerturbation, InputTransform
from botorch.models.transforms.outcome import OutcomeTransform
from botorch.utils.datasets import FixedNoiseDataset, RankingDataset, SupervisedDataset
from gpytorch.kernels import Kernel
from gpytorch.likelihoods.likelihood import Likelihood
from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
from torch import Tensor
NOT_YET_FIT_MSG = (
"Underlying BoTorch `Model` has not yet received its training_data. "
"Please fit the model first."
)
logger: Logger = get_logger(__name__)
[docs]class Surrogate(Base):
"""
**All classes in 'botorch_modular' directory are under
construction, incomplete, and should be treated as alpha
versions only.**
Ax wrapper for BoTorch ``Model``, subcomponent of ``BoTorchModel``
and is not meant to be used outside of it.
Args:
botorch_model_class: ``Model`` class to be used as the underlying
BoTorch model.
model_options: Dictionary of options / kwargs for the BoTorch
``Model`` constructed during ``Surrogate.fit``.
mll_class: ``MarginalLogLikelihood`` class to use for model-fitting.
mll_options: Dictionary of options / kwargs for the MLL.
outcome_transform: BoTorch outcome transforms. Passed down to the
BoTorch ``Model``. Multiple outcome transforms can be chained
together using ``ChainedOutcomeTransform``.
input_transform: BoTorch input transforms. Passed down to the
BoTorch ``Model``. Multiple input transforms can be chained
together using ``ChainedInputTransform``.
covar_module_class: Covariance module class, not yet used. Will be
used to construct custom BoTorch ``Model`` in the future.
covar_module_options: Covariance module kwargs, not yet used. Will be
used to construct custom BoTorch ``Model`` in the future.
likelihood: ``Likelihood`` class, not yet used. Will be used to
construct custom BoTorch ``Model`` in the future.
likelihood_options: Likelihood options, not yet used. Will be used to
construct custom BoTorch ``Model`` in the future.
"""
botorch_model_class: Type[Model]
model_options: Dict[str, Any]
mll_class: Type[MarginalLogLikelihood]
mll_options: Dict[str, Any]
outcome_transform: Optional[OutcomeTransform] = None
input_transform: Optional[InputTransform] = None
covar_module_class: Optional[Type[Kernel]] = None
covar_module_options: Dict[str, Any]
likelihood_class: Optional[Type[Likelihood]] = None
likelihood_options: Dict[str, Any]
_training_data: Optional[List[SupervisedDataset]] = None
_outcomes: Optional[List[str]] = None
_model: Optional[Model] = None
# Special setting for surrogates instantiated via `Surrogate.from_botorch`,
# to avoid re-constructing the underlying BoTorch model on `Surrogate.fit`
# when set to `False`.
_constructed_manually: bool = False
def __init__(
self,
# TODO: make optional when BoTorch model factory is checked in.
# Construction will then be possible from likelihood, kernel, etc.
botorch_model_class: Type[Model],
model_options: Optional[Dict[str, Any]] = None,
mll_class: Type[MarginalLogLikelihood] = ExactMarginalLogLikelihood,
mll_options: Optional[Dict[str, Any]] = None,
outcome_transform: Optional[OutcomeTransform] = None,
input_transform: Optional[InputTransform] = None,
covar_module_class: Optional[Type[Kernel]] = None,
covar_module_options: Optional[Dict[str, Any]] = None,
likelihood_class: Optional[Type[Likelihood]] = None,
likelihood_options: Optional[Dict[str, Any]] = None,
) -> None:
self.botorch_model_class = botorch_model_class
self.model_options = model_options or {}
self.mll_class = mll_class
self.mll_options = mll_options or {}
self.outcome_transform = outcome_transform
self.input_transform = input_transform
self.covar_module_class = covar_module_class
self.covar_module_options = covar_module_options or {}
self.likelihood_class = likelihood_class
self.likelihood_options = likelihood_options or {}
@property
def model(self) -> Model:
if self._model is None:
raise ValueError(
"BoTorch `Model` has not yet been constructed, please fit the "
"surrogate first (done via `BoTorchModel.fit`)."
)
return not_none(self._model)
@property
def training_data(self) -> List[SupervisedDataset]:
if self._training_data is None:
raise ValueError(NOT_YET_FIT_MSG)
return not_none(self._training_data)
@property
def Xs(self) -> List[Tensor]:
# Handles multi-output models. TODO: Improve this!
training_data = self.training_data
Xs = []
for dataset in training_data:
if self.botorch_model_class == PairwiseGP and isinstance(
dataset, RankingDataset
):
Xi = dataset.X.values
else:
Xi = dataset.X()
for _ in range(dataset.Y.shape[-1]):
Xs.append(Xi)
return Xs
@property
def dtype(self) -> torch.dtype:
return self.training_data[0].X.dtype
@property
def device(self) -> torch.device:
return self.training_data[0].X.device
[docs] @classmethod
def from_botorch(
cls,
model: Model,
mll_class: Type[MarginalLogLikelihood] = ExactMarginalLogLikelihood,
) -> Surrogate:
"""Instantiate a `Surrogate` from a pre-instantiated Botorch `Model`."""
surrogate = cls(botorch_model_class=model.__class__, mll_class=mll_class)
surrogate._model = model
# Temporarily disallowing `update` for surrogates instantiated from
# pre-made BoTorch `Model` instances to avoid reconstructing models
# that were likely pre-constructed for a reason (e.g. if this setup
# doesn't fully allow to constuct them).
surrogate._constructed_manually = True
return surrogate
[docs] def clone_reset(self) -> Surrogate:
return self.__class__(**self._serialize_attributes_as_kwargs())
[docs] def construct(self, datasets: List[SupervisedDataset], **kwargs: Any) -> None:
"""Constructs the underlying BoTorch ``Model`` using the training data.
Args:
training_data: Training data for the model (for one outcome for
the default `Surrogate`, with the exception of batched
multi-output case, where training data is formatted with just
one X and concatenated Ys).
**kwargs: Optional keyword arguments, expects any of:
- "fidelity_features": Indices of columns in X that represent
fidelity.
"""
if self._constructed_manually:
logger.warning("Reconstructing a manually constructed `Model`.")
if not len(datasets) == 1:
raise ValueError( # pragma: no cover
"Base `Surrogate` expects training data for single outcome."
)
input_constructor_kwargs = {**self.model_options, **(kwargs or {})}
dataset = datasets[0]
botorch_model_class_args = inspect.getfullargspec(self.botorch_model_class).args
# Temporary workaround to allow models to consume data from
# `FixedNoiseDataset`s even if they don't accept variance observations
if "train_Yvar" not in botorch_model_class_args and isinstance(
dataset, FixedNoiseDataset
):
warnings.warn(
f"Provided model class {self.botorch_model_class} does not accept "
"`train_Yvar` argument, but received `FixedNoiseDataset`. Ignoring "
"variance observations and converting to `SupervisedDataset`.",
AxWarning,
)
dataset = SupervisedDataset(X=dataset.X(), Y=dataset.Y())
self._training_data = [dataset]
formatted_model_inputs = self.botorch_model_class.construct_inputs(
training_data=dataset, **input_constructor_kwargs
)
self._set_formatted_inputs(
formatted_model_inputs=formatted_model_inputs,
inputs=[
[
"covar_module",
self.covar_module_class,
self.covar_module_options,
None,
],
["likelihood", self.likelihood_class, self.likelihood_options, None],
["outcome_transform", None, None, self.outcome_transform],
["input_transform", None, None, self.input_transform],
],
dataset=dataset,
botorch_model_class_args=botorch_model_class_args,
robust_digest=kwargs.get("robust_digest", None),
)
# pyre-ignore [45]
self._model = self.botorch_model_class(**formatted_model_inputs)
def _set_formatted_inputs(
self,
formatted_model_inputs: Dict[str, Any],
# pyre-fixme[2]: Parameter annotation cannot contain `Any`.
inputs: List[List[Any]],
dataset: SupervisedDataset,
# pyre-fixme[2]: Parameter annotation cannot be `Any`.
botorch_model_class_args: Any,
robust_digest: Optional[Dict[str, Any]] = None,
) -> None:
for input_name, input_class, input_options, input_object in inputs:
if input_class is None and input_object is None:
continue
if input_name not in botorch_model_class_args:
# TODO: We currently only pass in `covar_module` and `likelihood`
# if they are inputs to the BoTorch model. This interface will need
# to be expanded to a ModelFactory, see D22457664, to accommodate
# different models in the future.
raise UserInputError(
f"The BoTorch model class {self.botorch_model_class} does not "
f"support the input {input_name}."
)
if input_class is not None and input_object is not None: # pragma: no cover
raise RuntimeError(f"Got both a class and an object for {input_name}.")
if input_class is not None:
input_options = input_options or {}
formatted_model_inputs[input_name] = input_class(**input_options)
else:
formatted_model_inputs[input_name] = input_object
# Construct input perturbation if doing robust optimization.
if robust_digest is not None:
if len(robust_digest["environmental_variables"]):
# TODO[T131759269]: support env variables.
raise NotImplementedError(
"Environmental variable support is not yet implemented."
)
samples = torch.as_tensor(
robust_digest["sample_param_perturbations"](),
dtype=self.dtype,
device=self.device,
)
perturbation = InputPerturbation(
perturbation_set=samples, multiplicative=robust_digest["multiplicative"]
)
if formatted_model_inputs.get("input_transform") is not None:
# TODO: Support mixing with user supplied transforms.
raise NotImplementedError(
"User supplied input transforms are not supported "
"in robust optimization."
)
else:
formatted_model_inputs["input_transform"] = perturbation
[docs] def fit(
self,
datasets: List[SupervisedDataset],
metric_names: List[str],
search_space_digest: SearchSpaceDigest,
candidate_metadata: Optional[List[List[TCandidateMetadata]]] = None,
state_dict: Optional[Dict[str, Tensor]] = None,
refit: bool = True,
original_metric_names: Optional[List[str]] = None,
) -> None:
"""Fits the underlying BoTorch ``Model`` to ``m`` outcomes.
NOTE: ``state_dict`` and ``refit`` keyword arguments control how the
undelying BoTorch ``Model`` will be fit: whether its parameters will
be reoptimized and whether it will be warm-started from a given state.
There are three possibilities:
* ``fit(state_dict=None)``: fit model from scratch (optimize model
parameters and set its training data used for inference),
* ``fit(state_dict=some_state_dict, refit=True)``: warm-start refit
with a state dict of parameters (still re-optimize model parameters
and set the training data),
* ``fit(state_dict=some_state_dict, refit=False)``: load model parameters
without refitting, but set new training data (used in cross-validation,
for example).
Args:
datasets: A list of ``SupervisedDataset`` containers, each
corresponding to the data of one metric (outcome), to be passed
to ``Model.construct_inputs`` in BoTorch.
metric_names: A list of metric names, with the i-th metric
corresponding to the i-th dataset.
search_space_digest: A ``SearchSpaceDigest`` object containing
metadata on the features in the datasets.
candidate_metadata: Model-produced metadata for candidates, in
the order corresponding to the Xs.
state_dict: Optional state dict to load.
refit: Whether to re-optimize model parameters.
# TODO: we should refactor the fit() API to get rid of the metric_names
# and the concatenation hack that comes with it in BoTorchModel.fit()
# by attaching the individual metric_name to each dataset directly.
original_metric_names: sometimes the original list of metric_names
got tranformed into a different format before being passed down
into fit(). This arg preserves the original metric_names before
the transformation.
"""
if self._constructed_manually:
logger.debug(
"For manually constructed surrogates (via `Surrogate.from_botorch`), "
"`fit` skips setting the training data on model and only reoptimizes "
"its parameters if `refit=True`."
)
else:
self.construct(
datasets=datasets,
metric_names=metric_names,
**dataclasses.asdict(search_space_digest),
)
self._outcomes = (
original_metric_names
if original_metric_names is not None
else metric_names
)
if state_dict:
self.model.load_state_dict(not_none(state_dict))
if state_dict is None or refit:
fit_botorch_model(
model=self.model, mll_class=self.mll_class, mll_options=self.mll_options
)
[docs] def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
"""Predicts outcomes given a model and input tensor.
Args:
model: A botorch Model.
X: A ``n x d`` tensor of input parameters.
Returns:
Tensor: The predicted posterior mean as an ``n x o``-dim tensor.
Tensor: The predicted posterior covariance as a ``n x o x o``-dim tensor.
"""
# This temporarily disables the one-to-many transforms to avoid perturbing
# the user supplied parameterization.
with disable_one_to_many_transforms(self.model):
return predict_from_model(model=self.model, X=X)
[docs] def best_in_sample_point(
self,
search_space_digest: SearchSpaceDigest,
torch_opt_config: TorchOptConfig,
options: Optional[TConfig] = None,
) -> Tuple[Tensor, float]:
"""Finds the best observed point and the corresponding observed outcome
values.
"""
if torch_opt_config.is_moo:
raise NotImplementedError(
"Best observed point is incompatible with MOO problems."
)
best_point_and_observed_value = best_in_sample_point(
Xs=self.Xs,
# pyre-ignore[6]: `best_in_sample_point` currently expects a `TorchModel`
# as `model` kwarg, but only uses them for `predict` function, the
# signature for which is the same on this `Surrogate`.
# TODO: When we move `botorch_modular` directory to OSS, we will extend
# the annotation for `model` kwarg to accept `Surrogate` too.
model=self,
bounds=search_space_digest.bounds,
objective_weights=torch_opt_config.objective_weights,
outcome_constraints=torch_opt_config.outcome_constraints,
linear_constraints=torch_opt_config.linear_constraints,
fixed_features=torch_opt_config.fixed_features,
risk_measure=torch_opt_config.risk_measure,
options=options,
)
if best_point_and_observed_value is None:
raise ValueError("Could not obtain best in-sample point.")
best_point, observed_value = best_point_and_observed_value
return (
best_point.to(dtype=self.dtype, device=torch.device("cpu")),
observed_value,
)
[docs] def best_out_of_sample_point(
self,
search_space_digest: SearchSpaceDigest,
torch_opt_config: TorchOptConfig,
options: Optional[TConfig] = None,
) -> Tuple[Tensor, Tensor]:
"""Finds the best predicted point and the corresponding value of the
appropriate best point acquisition function.
"""
if torch_opt_config.fixed_features:
# When have fixed features, need `FixedFeatureAcquisitionFunction`
# which has peculiar instantiation (wraps another acquisition fn.),
# so need to figure out how to handle.
# TODO (ref: https://fburl.com/diff/uneqb3n9)
raise NotImplementedError("Fixed features not yet supported.")
options = options or {}
acqf_class, acqf_options = pick_best_out_of_sample_point_acqf_class(
outcome_constraints=torch_opt_config.outcome_constraints,
seed_inner=checked_cast_optional(int, options.get(Keys.SEED_INNER, None)),
qmc=checked_cast(bool, options.get(Keys.QMC, True)),
risk_measure=torch_opt_config.risk_measure,
)
# Avoiding circular import between `Surrogate` and `Acquisition`.
from ax.models.torch.botorch_modular.acquisition import Acquisition
acqf = Acquisition( # TODO: For multi-fidelity, might need diff. class.
surrogate=self,
botorch_acqf_class=acqf_class,
search_space_digest=search_space_digest,
torch_opt_config=torch_opt_config,
options=acqf_options,
)
candidates, acqf_values = acqf.optimize(
n=1,
search_space_digest=search_space_digest,
inequality_constraints=_to_inequality_constraints(
linear_constraints=torch_opt_config.linear_constraints
),
fixed_features=torch_opt_config.fixed_features,
)
return candidates[0], acqf_values[0]
[docs] def update(
self,
datasets: List[SupervisedDataset],
metric_names: List[str],
search_space_digest: SearchSpaceDigest,
candidate_metadata: Optional[List[List[TCandidateMetadata]]] = None,
state_dict: Optional[Dict[str, Tensor]] = None,
refit: bool = True,
) -> None:
"""Updates the surrogate model with new data. In the base ``Surrogate``,
just calls ``fit`` after checking that this surrogate was not created
via ``Surrogate.from_botorch`` (in which case the ``Model`` comes premade,
constructed manually and then supplied to ``Surrogate``).
NOTE: Expects `training_data` to be all available data,
not just the new data since the last time the model was updated.
Args:
training_data: Surrogate training_data containing all the data the model
should use for inference.
search_space_digest: A SearchSpaceDigest object containing
metadata on the features in the training data.
metric_names: Names of each outcome Y in Ys.
candidate_metadata: Model-produced metadata for candidates, in
the order corresponding to the Xs.
state_dict: Optional state dict to load.
refit: Whether to re-optimize model parameters or just set the training
data used for interence to new training data.
"""
# NOTE: In the future, could have `incremental` kwarg, in which case
# `training_data` could contain just the new data.
if self._constructed_manually:
raise NotImplementedError(
"`update` not yet implemented for models that are "
"constructed manually, but it is possible to create a new "
"surrogate in the same way as the current manually constructed one, "
"via `Surrogate.from_botorch`."
)
self.fit(
datasets=datasets,
metric_names=metric_names,
search_space_digest=search_space_digest,
candidate_metadata=candidate_metadata,
state_dict=state_dict,
refit=refit,
)
[docs] def pareto_frontier(self) -> Tuple[Tensor, Tensor]:
"""For multi-objective optimization, retrieve Pareto frontier instead
of best point.
Returns: A two-tuple of:
- tensor of points in the feature space,
- tensor of corresponding (multiple) outcomes.
"""
raise NotImplementedError(
"Pareto frontier not yet implemented."
) # pragma: no cover
[docs] def compute_diagnostics(self) -> Dict[str, Any]:
"""Computes model diagnostics like cross-validation measure of fit, etc."""
return {} # pragma: no cover
def _serialize_attributes_as_kwargs(self) -> Dict[str, Any]:
"""Serialize attributes of this surrogate, to be passed back to it
as kwargs on reinstantiation.
"""
if self._constructed_manually:
raise UnsupportedError(
"Surrogates constructed manually (ie Surrogate.from_botorch) may not "
"be serialized. If serialization is necessary please initialize from "
"the constructor."
)
return {
"botorch_model_class": self.botorch_model_class,
"model_options": self.model_options,
"mll_class": self.mll_class,
"mll_options": self.mll_options,
"outcome_transform": self.outcome_transform,
"input_transform": self.input_transform,
"covar_module_class": self.covar_module_class,
"covar_module_options": self.covar_module_options,
"likelihood_class": self.likelihood_class,
"likelihood_options": self.likelihood_options,
}
