Support for outcome transforms that return a TransformedPosterior in ModelListGP (#1563)

Summary:
Pull Request resolved: #1563

Replaces D41860896, but adds support for the outcome transforms that weren't working for `ModelListGP.posterior` by calling `ModelList.posterior`.

See #1519 for more context on the issue that this is fixing.

Reviewed By: saitcakmak

Differential Revision: D42019721

fbshipit-source-id: f2f566c53f327a02e26008428a187a6b9abf0c90

Author: esantorella

botorch/models/gpytorch.py

@@ -39,6 +39,7 @@
 from torch import Tensor
 
 if TYPE_CHECKING:
+    from botorch.posteriors.posterior_list import PosteriorList  # pragma: no cover
     from botorch.posteriors.transformed import TransformedPosterior  # pragma: no cover
     from gpytorch.likelihoods import Likelihood  # pragma: no cover
 
@@ -555,14 +556,15 @@ def batch_shape(self) -> torch.Size:
                 raise NotImplementedError(msg + " that are not broadcastble.")
         return next(iter(batch_shapes))
 
+    # pyre-fixme[15]: Inconsistent override in return types
     def posterior(
         self,
         X: Tensor,
         output_indices: Optional[List[int]] = None,
         observation_noise: Union[bool, Tensor] = False,
         posterior_transform: Optional[PosteriorTransform] = None,
         **kwargs: Any,
-    ) -> GPyTorchPosterior:
+    ) -> Union[GPyTorchPosterior, PosteriorList]:
         r"""Computes the posterior over model outputs at the provided points.
 
         Args:
@@ -582,11 +584,38 @@ def posterior(
             posterior_transform: An optional PosteriorTransform.
 
         Returns:
-            A `GPyTorchPosterior` or `FullyBayesianPosterior` object, representing
-            `batch_shape` joint distributions over `q` points and the outputs selected
-            by `output_indices` each. Includes measurement noise if
-            `observation_noise` is specified.
+            - If no `posterior_transform` is provided and the component models have no
+                `outcome_transform`, or if the component models only use linear outcome
+                transforms like `Standardize` (i.e. not `Log`), returns a
+                `GPyTorchPosterior` or `FullyBayesianPosterior` object,
+                representing `batch_shape` joint distributions over `q` points
+                and the outputs selected by `output_indices` each. Includes
+                measurement noise if `observation_noise` is specified.
+            - If no `posterior_transform` is provided and component models have
+                nonlinear transforms like `Log`, returns a `PosteriorList` with
+                sub-posteriors of type `TransformedPosterior`
+            - If `posterior_transform` is provided, that posterior transform will be
+               applied and will determine the return type. This could potentially be
+               any subclass of `Posterior`, but common choices give a
+               `GPyTorchPosterior`.
         """
+
+        # Nonlinear transforms untransform to a `TransformedPosterior`,
+        # which can't be made into a `GPyTorchPosterior`
+        returns_untransformed = any(
+            hasattr(mod, "outcome_transform") and (not mod.outcome_transform._is_linear)
+            for mod in self.models
+        )
+        if returns_untransformed:
+            return ModelList.posterior(
+                self,
+                X,
+                output_indices,
+                observation_noise,
+                posterior_transform,
+                **kwargs,
+            )
+
         self.eval()  # make sure model is in eval mode
         # input transforms are applied at `posterior` in `eval` mode, and at
         # `model.forward()` at the training time
@@ -628,10 +657,10 @@ def posterior(
         # apply output transforms of individual models if present
         mvns = []
         for i, mvn in mvn_gen:
-            try:
+            if hasattr(self.models[i], "outcome_transform"):
                 oct = self.models[i].outcome_transform
                 tf_mvn = oct.untransform_posterior(GPyTorchPosterior(mvn)).distribution
-            except AttributeError:
+            else:
                 tf_mvn = mvn
             mvns.append(tf_mvn)
         # return result as a GPyTorchPosteriors/FullyBayesianPosterior

botorch/models/model.py

@@ -391,7 +391,7 @@ def posterior(
         X: Tensor,
         output_indices: Optional[List[int]] = None,
         observation_noise: bool = False,
-        posterior_transform: Optional[Callable[[Posterior], Posterior]] = None,
+        posterior_transform: Optional[Callable[[PosteriorList], Posterior]] = None,
         **kwargs: Any,
     ) -> Posterior:
         r"""Computes the posterior over model outputs at the provided points.

botorch/models/transforms/outcome.py

@@ -101,8 +101,22 @@ def untransform(
             f"{self.__class__.__name__} does not implement the `untransform` method"
         )
 
+    @property
+    def _is_linear(self) -> bool:
+        """
+        True for transformations such as `Standardize`; these should be able to apply
+        `untransform_posterior` to a GPyTorchPosterior and return a GPyTorchPosterior,
+        because a multivariate normal distribution should remain multivariate normal
+        after applying the transform.
+        """
+        return False
+
     def untransform_posterior(self, posterior: Posterior) -> Posterior:
-        r"""Un-transform a posterior
+        r"""Un-transform a posterior.
+
+        Posteriors with `_is_linear=True` should return a `GPyTorchPosterior` when
+        `posterior` is a `GPyTorchPosterior`. Posteriors with `_is_linear=False`
+        likely return a `TransformedPosterior` instead.
 
         Args:
             posterior: A posterior in the transformed space.
@@ -182,6 +196,14 @@ def untransform(
             Y, Yvar = tf.untransform(Y, Yvar)
         return Y, Yvar
 
+    @property
+    def _is_linear(self) -> bool:
+        """
+        A `ChainedOutcomeTransform` is linear only if all of the component transforms
+        are linear.
+        """
+        return all((octf._is_linear for octf in self.values()))
+
     def untransform_posterior(self, posterior: Posterior) -> Posterior:
         r"""Un-transform a posterior
 
@@ -255,7 +277,10 @@ def forward(
             if Y.shape[:-2] != self._batch_shape:
                 raise RuntimeError("wrong batch shape")
             if Y.size(-1) != self._m:
-                raise RuntimeError("wrong output dimension")
+                raise RuntimeError(
+                    f"Wrong output dimension. Y.size(-1) is {Y.size(-1)}; expected "
+                    f"{self._m}."
+                )
             stdvs = Y.std(dim=-2, keepdim=True)
             stdvs = stdvs.where(stdvs >= self._min_stdv, torch.full_like(stdvs, 1.0))
             means = Y.mean(dim=-2, keepdim=True)
@@ -331,6 +356,10 @@ def untransform(
         Yvar_utf = self._stdvs_sq * Yvar if Yvar is not None else None
         return Y_utf, Yvar_utf
 
+    @property
+    def _is_linear(self) -> bool:
+        return True
+
     def untransform_posterior(
         self, posterior: Posterior
     ) -> Union[GPyTorchPosterior, TransformedPosterior]:

botorch/posteriors/gpytorch.py

@@ -35,6 +35,7 @@
 
 class GPyTorchPosterior(TorchPosterior):
     r"""A posterior based on GPyTorch's multi-variate Normal distributions."""
+    distribution: MultivariateNormal
 
     def __init__(
         self,

test/models/test_model_list_gp_regression.py

@@ -14,9 +14,9 @@
 from botorch.fit import fit_gpytorch_mll
 from botorch.models import ModelListGP
 from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
-from botorch.models.transforms import Standardize
 from botorch.models.transforms.input import Normalize
-from botorch.posteriors import GPyTorchPosterior
+from botorch.models.transforms.outcome import ChainedOutcomeTransform, Log, Standardize
+from botorch.posteriors import GPyTorchPosterior, PosteriorList, TransformedPosterior
 from botorch.sampling.normal import IIDNormalSampler
 from botorch.utils.testing import _get_random_data, BotorchTestCase
 from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
@@ -28,14 +28,34 @@
 from gpytorch.priors import GammaPrior
 
 
-def _get_model(fixed_noise=False, use_octf=False, use_intf=False, **tkwargs):
+def _get_model(
+    fixed_noise=False, outcome_transform: str = "None", use_intf=False, **tkwargs
+) -> ModelListGP:
     train_x1, train_y1 = _get_random_data(
         batch_shape=torch.Size(), m=1, n=10, **tkwargs
     )
+    train_y1 = torch.exp(train_y1)
     train_x2, train_y2 = _get_random_data(
         batch_shape=torch.Size(), m=1, n=11, **tkwargs
     )
-    octfs = [Standardize(m=1), Standardize(m=1)] if use_octf else [None, None]
+    if outcome_transform == "Standardize":
+        octfs = [Standardize(m=1), Standardize(m=1)]
+    elif outcome_transform == "Log":
+        octfs = [Log(), Standardize(m=1)]
+    elif outcome_transform == "Chained":
+        octfs = [
+            ChainedOutcomeTransform(
+                chained=ChainedOutcomeTransform(log=Log(), standardize=Standardize(m=1))
+            ),
+            Standardize(m=1),
+        ]
+    elif outcome_transform == "None":
+        octfs = [None, None]
+    else:
+        raise KeyError(  # pragma: no cover
+            "outcome_transform must be one of 'Standardize', 'Log', 'Chained', or "
+            "'None'."
+        )
     intfs = [Normalize(d=1), Normalize(d=1)] if use_intf else [None, None]
     if fixed_noise:
         train_y1_var = 0.1 + 0.1 * torch.rand_like(train_y1, **tkwargs)
@@ -73,10 +93,12 @@ def _get_model(fixed_noise=False, use_octf=False, use_intf=False, **tkwargs):
 
 class TestModelListGP(BotorchTestCase):
     def _base_test_ModelListGP(
-        self, fixed_noise: bool, dtype, use_octf: bool
+        self, fixed_noise: bool, dtype, outcome_transform: str
     ) -> ModelListGP:
         tkwargs = {"device": self.device, "dtype": dtype}
-        model = _get_model(fixed_noise=fixed_noise, use_octf=use_octf, **tkwargs)
+        model = _get_model(
+            fixed_noise=fixed_noise, outcome_transform=outcome_transform, **tkwargs
+        )
         self.assertIsInstance(model, ModelListGP)
         self.assertIsInstance(model.likelihood, LikelihoodList)
         for m in model.models:
@@ -85,8 +107,12 @@ def _base_test_ModelListGP(
             matern_kernel = m.covar_module.base_kernel
             self.assertIsInstance(matern_kernel, MaternKernel)
             self.assertIsInstance(matern_kernel.lengthscale_prior, GammaPrior)
-            if use_octf:
-                self.assertIsInstance(m.outcome_transform, Standardize)
+            if outcome_transform != "None":
+                self.assertIsInstance(
+                    m.outcome_transform, (Log, Standardize, ChainedOutcomeTransform)
+                )
+            else:
+                assert not hasattr(m, "outcome_transform")
 
         # test constructing likelihood wrapper
         mll = SumMarginalLogLikelihood(model.likelihood, model)
@@ -121,9 +147,19 @@ def _base_test_ModelListGP(
         # test posterior
         test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
         posterior = model.posterior(test_x)
-        self.assertIsInstance(posterior, GPyTorchPosterior)
-        self.assertIsInstance(posterior.distribution, MultitaskMultivariateNormal)
-        if use_octf:
+        gpytorch_posterior_expected = outcome_transform in ("None", "Standardize")
+        expected_type = (
+            GPyTorchPosterior if gpytorch_posterior_expected else PosteriorList
+        )
+        self.assertIsInstance(posterior, expected_type)
+        submodel = model.models[0]
+        p0 = submodel.posterior(test_x)
+        self.assertTrue(torch.allclose(posterior.mean[:, [0]], p0.mean))
+        self.assertTrue(torch.allclose(posterior.variance[:, [0]], p0.variance))
+
+        if gpytorch_posterior_expected:
+            self.assertIsInstance(posterior.distribution, MultitaskMultivariateNormal)
+        if outcome_transform != "None":
             # ensure un-transformation is applied
             submodel = model.models[0]
             p0 = submodel.posterior(test_x)
@@ -136,8 +172,9 @@ def _base_test_ModelListGP(
 
         # test output_indices
         posterior = model.posterior(test_x, output_indices=[0], observation_noise=True)
-        self.assertIsInstance(posterior, GPyTorchPosterior)
-        self.assertIsInstance(posterior.distribution, MultivariateNormal)
+        self.assertIsInstance(posterior, expected_type)
+        if gpytorch_posterior_expected:
+            self.assertIsInstance(posterior.distribution, MultivariateNormal)
 
         # test condition_on_observations
         f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
@@ -176,39 +213,50 @@ def _base_test_ModelListGP(
         X = torch.rand(3, 1, **tkwargs)
         weights = torch.tensor([1, 2], **tkwargs)
         post_tf = ScalarizedPosteriorTransform(weights=weights)
-        posterior_tf = model.posterior(X, posterior_transform=post_tf)
-        self.assertTrue(
-            torch.allclose(
-                posterior_tf.mean,
-                model.posterior(X).mean @ weights.unsqueeze(-1),
+        if gpytorch_posterior_expected:
+            posterior_tf = model.posterior(X, posterior_transform=post_tf)
+            self.assertTrue(
+                torch.allclose(
+                    posterior_tf.mean,
+                    model.posterior(X).mean @ weights.unsqueeze(-1),
+                )
             )
-        )
 
         return model
 
     def test_ModelListGP(self) -> None:
-        for dtype, use_octf in itertools.product(
-            (torch.float, torch.double), (False, True)
+        for dtype, outcome_transform in itertools.product(
+            (torch.float, torch.double), ("None", "Standardize", "Log", "Chained")
         ):
 
             model = self._base_test_ModelListGP(
-                fixed_noise=False, dtype=dtype, use_octf=use_octf
+                fixed_noise=False, dtype=dtype, outcome_transform=outcome_transform
             )
             tkwargs = {"device": self.device, "dtype": dtype}
 
             # test observation_noise
             test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
             posterior = model.posterior(test_x, observation_noise=True)
-            self.assertIsInstance(posterior, GPyTorchPosterior)
-            self.assertIsInstance(posterior.distribution, MultitaskMultivariateNormal)
+
+            gpytorch_posterior_expected = outcome_transform in ("None", "Standardize")
+            expected_type = (
+                GPyTorchPosterior if gpytorch_posterior_expected else PosteriorList
+            )
+            self.assertIsInstance(posterior, expected_type)
+            if gpytorch_posterior_expected:
+                self.assertIsInstance(
+                    posterior.distribution, MultitaskMultivariateNormal
+                )
+            else:
+                self.assertIsInstance(posterior.posteriors[0], TransformedPosterior)
 
     def test_ModelListGP_fixed_noise(self) -> None:
 
-        for dtype, use_octf in itertools.product(
-            (torch.float, torch.double), (False, True)
+        for dtype, outcome_transform in itertools.product(
+            (torch.float, torch.double), ("None", "Standardize")
         ):
             model = self._base_test_ModelListGP(
-                fixed_noise=True, dtype=dtype, use_octf=use_octf
+                fixed_noise=True, dtype=dtype, outcome_transform=outcome_transform
             )
             tkwargs = {"device": self.device, "dtype": dtype}
             f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]

test/models/transforms/test_outcome.py

@@ -93,10 +93,34 @@ def test_standardize_raises_when_mean_not_set(self) -> None:
             ):
                 transform.untransform(y)
 
+    def test_is_linear(self) -> None:
+        posterior = _get_test_posterior(
+            shape=torch.Size([1, 1]), device=self.device, dtype=torch.float64
+        )
+        y = torch.arange(2, dtype=torch.float64, device=self.device)[:, None]
+        standardize_tf = Standardize(m=1)
+        standardize_tf(y)
+
+        for transform in [
+            standardize_tf,
+            Power(power=0.5),
+            Log(),
+            ChainedOutcomeTransform(
+                chained=ChainedOutcomeTransform(stand=standardize_tf)
+            ),
+            ChainedOutcomeTransform(log=Log()),
+        ]:
+            posterior_is_gpt = isinstance(
+                transform.untransform_posterior(posterior), GPyTorchPosterior
+            )
+            self.assertEqual(posterior_is_gpt, transform._is_linear)
+
     def test_standardize(self):
         # test error on incompatible dim
         tf = Standardize(m=1)
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(
+            RuntimeError, msg="Wrong output dimension. Y.size(-1) is 2; expected 1."
+        ):
             tf(torch.zeros(3, 2, device=self.device), None)
         # test error on incompatible batch shape
         with self.assertRaises(RuntimeError):