diff --git a/botorch/acquisition/multi_objective/utils.py b/botorch/acquisition/multi_objective/utils.py index 99197fa5eb..bb4502b137 100644 --- a/botorch/acquisition/multi_objective/utils.py +++ b/botorch/acquisition/multi_objective/utils.py @@ -27,7 +27,7 @@ from botorch.models.fully_bayesian import MCMC_DIM from botorch.models.model import Model from botorch.sampling.get_sampler import get_sampler -from botorch.utils.gp_sampling import get_gp_samples +from botorch.sampling.pathwise.posterior_samplers import get_matheron_path_model from botorch.utils.multi_objective.box_decompositions.box_decomposition import ( BoxDecomposition, ) @@ -320,7 +320,6 @@ def sample_optimal_points( optimizer: Callable[ [GenericDeterministicModel, Tensor, int, bool, Any], Tuple[Tensor, Tensor] ] = random_search_optimizer, - num_rff_features: int = 512, maximize: bool = True, optimizer_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[Tensor, Tensor]: @@ -344,7 +343,6 @@ def sample_optimal_points( num_samples: The number of GP samples. num_points: The number of optimal points to be outputted. optimizer: A callable that solves the deterministic optimization problem. - num_rff_features: The number of random Fourier features. maximize: If true, we consider a maximization problem. optimizer_kwargs: The additional arguments for the optimizer. @@ -356,7 +354,7 @@ def sample_optimal_points( - A `num_samples x num_points x M`-dim Tensor containing the collection of optimal objectives. """ - tkwargs = {"dtype": bounds.dtype, "device": bounds.device} + tkwargs: Dict[str, Any] = {"dtype": bounds.dtype, "device": bounds.device} M = model.num_outputs d = bounds.shape[-1] if M == 1: @@ -369,9 +367,7 @@ def sample_optimal_points( pareto_sets = torch.zeros((num_samples, num_points, d), **tkwargs) pareto_fronts = torch.zeros((num_samples, num_points, M), **tkwargs) for i in range(num_samples): - sample_i = get_gp_samples( - model=model, num_outputs=M, n_samples=1, num_rff_features=num_rff_features - ) + sample_i = get_matheron_path_model(model=model) ps_i, pf_i = optimizer( model=sample_i, bounds=bounds, diff --git a/botorch/models/deterministic.py b/botorch/models/deterministic.py index 96b4fc48c6..2b878466c3 100644 --- a/botorch/models/deterministic.py +++ b/botorch/models/deterministic.py @@ -19,8 +19,8 @@ observation. `GenericDeterministicModel` supports arbitrary deterministic functions, while `AffineFidelityCostModel` is a particular cost model for multi-fidelity optimization. Other use cases of deterministic models include -representing approximate GP sample paths, e.g. random Fourier features obtained -with `get_gp_samples`, which allows them to be substituted in acquisition +representing approximate GP sample paths, e.g. Matheron paths obtained +with `get_matheron_path_model`, which allows them to be substituted in acquisition functions or in other places where a `Model` is expected. """ diff --git a/botorch/sampling/pathwise/__init__.py b/botorch/sampling/pathwise/__init__.py index b78b774b15..6554053636 100644 --- a/botorch/sampling/pathwise/__init__.py +++ b/botorch/sampling/pathwise/__init__.py @@ -18,6 +18,7 @@ ) from botorch.sampling.pathwise.posterior_samplers import ( draw_matheron_paths, + get_matheron_path_model, MatheronPath, ) from botorch.sampling.pathwise.prior_samplers import draw_kernel_feature_paths @@ -28,6 +29,7 @@ "draw_matheron_paths", "draw_kernel_feature_paths", "gen_kernel_features", + "get_matheron_path_model", "gaussian_update", "GeneralizedLinearPath", "KernelEvaluationMap", diff --git a/botorch/sampling/pathwise/posterior_samplers.py b/botorch/sampling/pathwise/posterior_samplers.py index 0872268ad0..2f68bc9988 100644 --- a/botorch/sampling/pathwise/posterior_samplers.py +++ b/botorch/sampling/pathwise/posterior_samplers.py @@ -19,7 +19,10 @@ from typing import Optional, Union +import torch from botorch.models.approximate_gp import ApproximateGPyTorchModel +from botorch.models.deterministic import GenericDeterministicModel +from botorch.models.model import ModelList from botorch.models.model_list_gp_regression import ModelListGP from botorch.sampling.pathwise.paths import PathDict, PathList, SamplePath from botorch.sampling.pathwise.prior_samplers import ( @@ -36,8 +39,9 @@ ) from botorch.utils.context_managers import delattr_ctx from botorch.utils.dispatcher import Dispatcher +from botorch.utils.transforms import is_ensemble from gpytorch.models import ApproximateGP, ExactGP, GP -from torch import Size +from torch import Size, Tensor DrawMatheronPaths = Dispatcher("draw_matheron_paths") @@ -83,13 +87,66 @@ def __init__( ) +def get_matheron_path_model( + model: GP, sample_shape: Optional[Size] = None +) -> GenericDeterministicModel: + r"""Generates a deterministic model using a single Matheron path drawn + from the model's posterior. + + The deterministic model evalutes the output of `draw_matheron_paths`, + and reshapes it to mimic the output behavior of the model's posterior. + + Args: + model: The model whose posterior is to be sampled. + sample_shape: The shape of the sample paths to be drawn, if an ensemble + of sample paths is desired. If this is specified, the resulting + deterministic model will behave as if the `sample_shape` is prepended + to the `batch_shape` of the model. The inputs used to evaluate the model + must be adjusted to match. + + Returns: A deterministic model that evaluates the Matheron path. + """ + sample_shape = sample_shape or Size() + path = draw_matheron_paths(model, sample_shape=sample_shape) + num_outputs = model.num_outputs + + def f(X: Tensor) -> Tensor: + r"""Reshapes the path evaluations to bring the output dimension to the end. + + Args: + X: The input tensor of shape `batch_shape x q x d`. + If the model is batched, `batch_shape` must be broadcastable to + the model batch shape. + + Returns: + The output tensor of shape `batch_shape x q x m`. + """ + if num_outputs == 1: + # For single-output, we lack the output dimension. Add one. + res = path(X).unsqueeze(-1) + elif isinstance(model, ModelList): + # For model list, path evaluates to a list of tensors. Stack them. + res = torch.stack(path(X), dim=-1) + else: + # For multi-output, path expects inputs broadcastable to + # `model._aug_batch_shape x q x d` and returns outputs of shape + # `model._aug_batch_shape x q`. Augmented batch shape includes the + # `m` dimension, so we will unsqueeze that and transpose after. + res = path(X.unsqueeze(-3)).transpose(-1, -2) + return res + + path_model = GenericDeterministicModel(f=f, num_outputs=num_outputs) + path_model._is_ensemble = is_ensemble(model) or len(sample_shape) > 0 + return path_model + + def draw_matheron_paths( model: GP, sample_shape: Size, prior_sampler: TPathwisePriorSampler = draw_kernel_feature_paths, update_strategy: TPathwiseUpdate = gaussian_update, ) -> MatheronPath: - r"""Generates function draws from (an approximate) Gaussian process prior. + r"""Generates function draws from (an approximate) Gaussian process posterior. When evaluted, sample paths produced by this method return Tensors with dimensions `sample_dims x batch_dims x [joint_dim]`, where `joint_dim` denotes the penultimate diff --git a/botorch/utils/gp_sampling.py b/botorch/utils/gp_sampling.py index 2fdb885f87..008be0d2b1 100644 --- a/botorch/utils/gp_sampling.py +++ b/botorch/utils/gp_sampling.py @@ -6,6 +6,7 @@ from __future__ import annotations +import warnings from copy import deepcopy from math import pi from typing import List, Optional @@ -41,6 +42,14 @@ def __init__(self, model: Model, seed: Optional[int] = None) -> None: Args: model: The Model defining the GP prior. """ + warnings.warn( + "`GPDraw` is deprecated and will be removed in v0.13 release. " + "For drawing GP sample paths, we recommend using pathwise " + "sampling code found in `botorch/sampling/pathwise`. We recommend " + "`get_matheron_path_model` for most use cases.", + DeprecationWarning, + stacklevel=2, + ) super().__init__() self._model = deepcopy(model) self._num_outputs = self._model.num_outputs @@ -429,6 +438,14 @@ def get_gp_samples( A `GenericDeterministicModel` that evaluates `n_samples` sampled functions. If `n_samples > 1`, this will be a batched model. """ + warnings.warn( + "`get_gp_samples` is deprecated and will be removed in v0.13 release. " + "For drawing GP sample paths, we recommend using pathwise " + "sampling code found in `botorch/sampling/pathwise`. We recommend " + "`get_matheron_path_model` for most use cases.", + DeprecationWarning, + stacklevel=2, + ) # Get transforms from the model. intf = getattr(model, "input_transform", None) octf = getattr(model, "outcome_transform", None) diff --git a/test/acquisition/multi_objective/test_utils.py b/test/acquisition/multi_objective/test_utils.py index 621e9230fb..acdfddbc95 100644 --- a/test/acquisition/multi_objective/test_utils.py +++ b/test/acquisition/multi_objective/test_utils.py @@ -24,7 +24,7 @@ from botorch.models.gp_regression import SingleTaskGP from botorch.models.model_list_gp_regression import ModelListGP from botorch.models.transforms.outcome import Standardize -from botorch.utils.gp_sampling import get_gp_samples +from botorch.sampling.pathwise import get_matheron_path_model from botorch.utils.multi_objective import is_non_dominated from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior from torch import Tensor @@ -306,11 +306,7 @@ def test_random_search_optimizer(self): **tkwargs, ) - model_sample = get_gp_samples( - model=model, - num_outputs=num_objectives, - n_samples=1, - ) + model_sample = get_matheron_path_model(model=model) input_dim = X.shape[-1] # fake bounds diff --git a/test/sampling/pathwise/test_posterior_samplers.py b/test/sampling/pathwise/test_posterior_samplers.py index 1d8a943064..e6c2a37370 100644 --- a/test/sampling/pathwise/test_posterior_samplers.py +++ b/test/sampling/pathwise/test_posterior_samplers.py @@ -7,12 +7,15 @@ from __future__ import annotations from copy import deepcopy +from typing import Any, Dict import torch from botorch.models import ModelListGP, SingleTaskGP, SingleTaskVariationalGP +from botorch.models.deterministic import GenericDeterministicModel from botorch.models.transforms.input import Normalize from botorch.models.transforms.outcome import Standardize from botorch.sampling.pathwise import draw_matheron_paths, MatheronPath, PathList +from botorch.sampling.pathwise.posterior_samplers import get_matheron_path_model from botorch.sampling.pathwise.utils import get_train_inputs from botorch.utils.test_helpers import get_sample_moments, standardize_moments from botorch.utils.testing import BotorchTestCase @@ -24,7 +27,7 @@ class TestPosteriorSamplers(BotorchTestCase): def setUp(self, suppress_input_warnings: bool = True) -> None: super().setUp(suppress_input_warnings=suppress_input_warnings) - tkwargs = {"device": self.device, "dtype": torch.float64} + tkwargs: Dict[str, Any] = {"device": self.device, "dtype": torch.float64} torch.manual_seed(0) base = MaternKernel(nu=2.5, ard_num_dims=2, batch_shape=Size([])) @@ -67,6 +70,8 @@ def setUp(self, suppress_input_warnings: bool = True) -> None: outcome_transform=outcome_transform, ).to(**tkwargs) + self.tkwargs = tkwargs + def test_draw_matheron_paths(self): for seed, model in enumerate( (self.inferred_noise_gp, self.observed_noise_gp, self.variational_gp) @@ -122,3 +127,53 @@ def _test_draw_matheron_paths(self, model, paths, sample_shape, atol=3): tol = atol * (num_features**-0.5 + sample_shape.numel() ** -0.5) for exact, estimate in zip(exact_moments, sample_moments): self.assertTrue(exact.allclose(estimate, atol=tol, rtol=0)) + + def test_get_matheron_path_model(self) -> None: + model_list = ModelListGP(self.inferred_noise_gp, self.observed_noise_gp) + moo_model = SingleTaskGP( + train_X=torch.rand(5, 2, **self.tkwargs), + train_Y=torch.rand(5, 2, **self.tkwargs), + ) + + test_X = torch.rand(5, 2, **self.tkwargs) + batch_test_X = torch.rand(3, 5, 2, **self.tkwargs) + sample_shape = Size([2]) + sample_shape_X = torch.rand(3, 2, 5, 2, **self.tkwargs) + for model in (self.inferred_noise_gp, moo_model, model_list): + path_model = get_matheron_path_model(model=model) + self.assertFalse(path_model._is_ensemble) + self.assertIsInstance(path_model, GenericDeterministicModel) + for X in (test_X, batch_test_X): + self.assertEqual( + model.posterior(X).mean.shape, path_model.posterior(X).mean.shape + ) + path_model = get_matheron_path_model(model=model, sample_shape=sample_shape) + self.assertTrue(path_model._is_ensemble) + self.assertEqual( + path_model.posterior(sample_shape_X).mean.shape, + sample_shape_X.shape[:-1] + Size([model.num_outputs]), + ) + + def test_get_matheron_path_model_batched(self) -> None: + model = SingleTaskGP( + train_X=torch.rand(4, 5, 2, **self.tkwargs), + train_Y=torch.rand(4, 5, 2, **self.tkwargs), + ) + model._is_ensemble = True + path_model = get_matheron_path_model(model=model) + self.assertTrue(path_model._is_ensemble) + test_X = torch.rand(5, 2, **self.tkwargs) + # This mimics the behavior of the acquisition functions unsqueezing the + # model batch dimension for ensemble models. + batch_test_X = torch.rand(3, 1, 5, 2, **self.tkwargs) + # Explicitly matching X for completeness. + complete_test_X = torch.rand(3, 4, 5, 2, **self.tkwargs) + for X in (test_X, batch_test_X, complete_test_X): + self.assertEqual( + model.posterior(X).mean.shape, path_model.posterior(X).mean.shape + ) + + # Test with sample_shape. + path_model = get_matheron_path_model(model=model, sample_shape=Size([2, 6])) + test_X = torch.rand(3, 2, 6, 4, 5, 2, **self.tkwargs) + self.assertEqual(path_model.posterior(test_X).mean.shape, test_X.shape) diff --git a/test/utils/test_gp_sampling.py b/test/utils/test_gp_sampling.py index 2bb6c3d521..191a9627ca 100644 --- a/test/utils/test_gp_sampling.py +++ b/test/utils/test_gp_sampling.py @@ -190,7 +190,10 @@ def test_gp_draw_single_output(self): tkwargs = {"device": self.device, "dtype": dtype} model, _, _ = _get_model(**tkwargs) mean = model.mean_module.raw_constant.detach().clone() - gp = GPDraw(model) + with self.assertWarnsRegex( + DeprecationWarning, "is deprecated and will be removed" + ): + gp = GPDraw(model) # test initialization self.assertIsNone(gp.Xs) self.assertIsNone(gp.Ys) @@ -547,16 +550,19 @@ def test_get_gp_samples(self): ) with torch.random.fork_rng(): torch.manual_seed(0) - gp_samples = get_gp_samples( - model=( - batched_to_model_list(model) - if ((not use_batch_model) and (m > 1)) - else model - ), - num_outputs=m, - n_samples=n_samples, - num_rff_features=512, - ) + with self.assertWarnsRegex( + DeprecationWarning, "is deprecated and will be removed" + ): + gp_samples = get_gp_samples( + model=( + batched_to_model_list(model) + if ((not use_batch_model) and (m > 1)) + else model + ), + num_outputs=m, + n_samples=n_samples, + num_rff_features=512, + ) samples = gp_samples.posterior(X).mean self.assertEqual(samples.shape[0], n_samples) if batched_inputs: