Add a helper for producing a DeterministicModel using a Matheron path (#2435)

Summary:
Pull Request resolved: #2435

Out of the box, pathwise sampling code does not produce outputs that conform to BoTorch modeling conventions. This adds a helper that produces a Matheron path and wraps it in a `GenericDeterministicModel` using a helper that reshapes the path outputs to mimic the expected behavior from the model's posterior.

This will offer a straightforward replacement of `get_gp_samples`, which is being deprecated.

Reviewed By: esantorella

Differential Revision: D59941730

fbshipit-source-id: 6d5fe5cc2a48257bb564719f6dd63567c2d9b466

Author: saitcakmak

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,11 @@
 
 from typing import Optional, Union
 
+import torch
+from botorch.exceptions.errors import UnsupportedError
 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 +40,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 +88,69 @@ 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 = Size() if sample_shape is None else sample_shape
+    path = draw_matheron_paths(model, sample_shape=sample_shape)
+    num_outputs = model.num_outputs
+    if isinstance(model, ModelList) and len(model.models) != num_outputs:
+        raise UnsupportedError("A model-list of multi-output models is not supported.")
+
+    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

test/sampling/pathwise/test_posterior_samplers.py

@@ -7,12 +7,16 @@
 from __future__ import annotations
 
 from copy import deepcopy
+from typing import Any, Dict
 
 import torch
+from botorch.exceptions.errors import UnsupportedError
 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 +28,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 +71,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 +128,60 @@ 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]),
+            )
+
+        with self.assertRaisesRegex(
+            UnsupportedError, "A model-list of multi-output models is not supported."
+        ):
+            get_matheron_path_model(
+                model=ModelListGP(self.inferred_noise_gp, moo_model)
+            )
+
+    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)