Deprecate gp_sampling module in favor of pathwise sampling

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,

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.
 """
 

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",

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

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)

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

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)

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: