Add outcome_transform to FixedNoiseMultiTaskGP (#1255)

Summary:
Pull Request resolved: #1255

Add outcome_transform to FixedNoiseMultiTaskGP

Reviewed By: qingfeng10

Differential Revision: D37020585

fbshipit-source-id: a2613c4bc34f8aa10018be799dd3ff6ec38dda08

Author: ItsMrLin

botorch/models/multitask.py

@@ -113,6 +113,10 @@ def __init__(
                 on p.s.d. matrices. A common prior for this is the `LKJ` prior.
             input_transform: An input transform that is applied in the model's
                 forward pass.
+            outcome_transform: An outcome transform that is applied to the
+                training data during instantiation and to the posterior during
+                inference (that is, the `Posterior` obtained by calling
+                `.posterior` on the model will be on the original scale).
 
         Example:
             >>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
@@ -308,6 +312,7 @@ def __init__(
         output_tasks: Optional[List[int]] = None,
         rank: Optional[int] = None,
         input_transform: Optional[InputTransform] = None,
+        outcome_transform: Optional[OutcomeTransform] = None,
     ) -> None:
         r"""Multi-Task GP model using an ICM kernel and known observation noise.
 
@@ -328,6 +333,10 @@ def __init__(
                 full rank (i.e. number of tasks) kernel.
             input_transform: An input transform that is applied in the model's
                 forward pass.
+            outcome_transform: An outcome transform that is applied to the
+                training data during instantiation and to the posterior during
+                inference (that is, the `Posterior` obtained by calling
+                `.posterior` on the model will be on the original scale).
 
         Example:
             >>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
@@ -344,6 +353,10 @@ def __init__(
                 X=train_X, input_transform=input_transform
             )
         self._validate_tensor_args(X=transformed_X, Y=train_Y, Yvar=train_Yvar)
+
+        if outcome_transform is not None:
+            train_Y, train_Yvar = outcome_transform(train_Y, train_Yvar)
+
         # We'll instatiate a MultiTaskGP and simply override the likelihood
         super().__init__(
             train_X=train_X,
@@ -354,7 +367,11 @@ def __init__(
             rank=rank,
             task_covar_prior=task_covar_prior,
             input_transform=input_transform,
+            outcome_transform=None,  # outcome_transform is applied already
         )
+
+        if outcome_transform is not None:
+            self.outcome_transform = outcome_transform
         self.likelihood = FixedNoiseGaussianLikelihood(noise=train_Yvar.squeeze(-1))
         self.to(train_X)
 

test/models/test_multitask.py

@@ -83,7 +83,7 @@ def _gen_model_single_output(**tkwargs):
 
 
 def _gen_fixed_noise_model_and_data(
-    task_feature: int = 0, input_transform=None, **tkwargs
+    task_feature: int = 0, input_transform=None, outcome_transform=None, **tkwargs
 ):
     datasets, (train_X, train_Y, train_Yvar) = _gen_datasets(yvar=0.05, **tkwargs)
     model = FixedNoiseMultiTaskGP(
@@ -92,6 +92,7 @@ def _gen_fixed_noise_model_and_data(
         train_Yvar,
         task_feature=task_feature,
         input_transform=input_transform,
+        outcome_transform=outcome_transform,
     )
     return model.to(**tkwargs), datasets, (train_X, train_Y, train_Yvar)
 
@@ -345,17 +346,18 @@ def test_MultiTaskGP_given_covar_module(self):
 class TestFixedNoiseMultiTaskGP(BotorchTestCase):
     def test_FixedNoiseMultiTaskGP(self):
         bounds = torch.tensor([[-1.0, 0.0], [1.0, 1.0]])
-        for dtype, use_intf in itertools.product(
-            (torch.float, torch.double), (False, True)
+        for dtype, use_intf, use_octf in itertools.product(
+            (torch.float, torch.double), (False, True), (False, True)
         ):
             tkwargs = {"device": self.device, "dtype": dtype}
+            octf = Standardize(m=1) if use_octf else None
             intf = (
                 Normalize(d=2, bounds=bounds.to(**tkwargs), transform_on_train=True)
                 if use_intf
                 else None
             )
             model, _, (train_X, _, _) = _gen_fixed_noise_model_and_data(
-                input_transform=intf, **tkwargs
+                input_transform=intf, outcome_transform=octf, **tkwargs
             )
             self.assertIsInstance(model, FixedNoiseMultiTaskGP)
             self.assertEqual(model.num_outputs, 2)
@@ -370,6 +372,8 @@ def test_FixedNoiseMultiTaskGP(self):
             self.assertEqual(
                 model.task_covar_module.covar_factor.shape[-1], model._rank
             )
+            if use_octf:
+                self.assertIsInstance(model.outcome_transform, Standardize)
             if use_intf:
                 self.assertIsInstance(model.input_transform, Normalize)
 
@@ -394,6 +398,16 @@ def test_FixedNoiseMultiTaskGP(self):
             self.assertEqual(posterior_f.mean.shape, torch.Size([2, 2]))
             self.assertEqual(posterior_f.variance.shape, torch.Size([2, 2]))
 
+            # check posterior transform is applied
+            if use_octf:
+                posterior_pred = model.posterior(test_x)
+                tmp_tf = model.outcome_transform
+                del model.outcome_transform
+                pp_tf = model.posterior(test_x)
+                model.outcome_transform = tmp_tf
+                expected_var = tmp_tf.untransform_posterior(pp_tf).variance
+                self.assertTrue(torch.allclose(posterior_pred.variance, expected_var))
+
             # test that posterior w/ observation noise raises appropriate error
             with self.assertRaises(NotImplementedError):
                 model.posterior(test_x, observation_noise=True)