Pulled out condition_on_observations-tests for easier testing (#2988)

Summary:

Pulls out the existing condition_on_observations-tests for easier testing of the specific method. Does not need to land, but simplifies testing of condition_on_observations methods.

Reviewed By: Balandat

Differential Revision: D80805810

Author: Carl Hvarfner

test/models/test_fully_bayesian.py

@@ -792,7 +792,7 @@ def test_construct_inputs(self) -> None:
             else:
                 self.assertTrue(Yvar.equal(data_dict["train_Yvar"]))
 
-    def test_condition_on_observation(self) -> None:
+    def test_fbstgp_condition_on_observations(self) -> None:
         # The following conditioned data shapes should work (output describes):
         # training data shape after cond(batch shape in output is req. in gpytorch)
         # X: num_models x n x d, Y: num_models x n x d --> num_models x n x d

test/models/test_fully_bayesian_multitask.py

@@ -682,7 +682,7 @@ def test_acquisition_functions(self):
                         )
                     self.assertEqual(acqf(test_X).shape, torch.Size(batch_shape))
 
-    def test_condition_on_observation(self) -> None:
+    def test_condition_on_observations(self) -> None:
         # The following conditioned data shapes should work (output describes):
         # training data shape after cond(batch shape in output is req. in gpytorch)
         # X: num_models x n x d, Y: num_models x n x d --> num_models x n x d

test/models/test_gpytorch.py

@@ -155,12 +155,6 @@ def test_gpytorch_model(self):
             # test noise shape validation
             with self.assertRaises(BotorchTensorDimensionError):
                 model.posterior(test_X, observation_noise=torch.rand(2, **tkwargs))
-            # test conditioning on observations
-            cm = model.condition_on_observations(
-                torch.rand(2, 1, **tkwargs), torch.rand(2, 1, **tkwargs)
-            )
-            self.assertIsInstance(cm, SimpleGPyTorchModel)
-            self.assertEqual(cm.train_targets.shape, torch.Size([7]))
             # test subset_output
             with self.assertRaises(NotImplementedError):
                 model.subset_output([0])
@@ -255,20 +249,6 @@ def test_validate_tensor_args(self) -> None:
                     ):
                         GPyTorchModel._validate_tensor_args(X, Y, Yvar, strict=strict)
 
-    def test_condition_on_observations_tensor_validation(self) -> None:
-        model = SimpleGPyTorchModel(torch.rand(5, 1), torch.randn(5, 1))
-        model.posterior(torch.rand(2, 1))  # evaluate the model to form caches.
-        # Outside of fantasize, the inputs are validated.
-        with self.assertWarnsRegex(
-            BotorchTensorDimensionWarning, "Non-strict enforcement of"
-        ):
-            model.condition_on_observations(torch.randn(2, 1), torch.randn(5, 2, 1))
-        # Inside of fantasize, the inputs are not validated.
-        with fantasize(), warnings.catch_warnings(record=True) as ws:
-            warnings.filterwarnings("always", category=BotorchTensorDimensionWarning)
-            model.condition_on_observations(torch.randn(2, 1), torch.randn(5, 2, 1))
-        self.assertFalse(any(w.category is BotorchTensorDimensionWarning for w in ws))
-
     def test_fantasize_flag(self):
         train_X = torch.rand(5, 1)
         train_Y = torch.sin(train_X)
@@ -358,12 +338,6 @@ def test_batched_multi_output_gpytorch_model(self):
             # test subset_output
             with self.assertRaises(NotImplementedError):
                 model.subset_output([0])
-            # test conditioning on observations
-            cm = model.condition_on_observations(
-                torch.rand(2, 1, **tkwargs), torch.rand(2, 2, **tkwargs)
-            )
-            self.assertIsInstance(cm, SimpleBatchedMultiOutputGPyTorchModel)
-            self.assertEqual(cm.train_targets.shape, torch.Size([2, 7]))
             # test fantasize
             sampler = SobolQMCNormalSampler(sample_shape=torch.Size([2]))
             cm = model.fantasize(torch.rand(2, 1, **tkwargs), sampler=sampler)
@@ -402,6 +376,56 @@ def test_batched_multi_output_gpytorch_model(self):
                         ),
                     )
 
+    def test_condition_on_observations(self):
+        for dtype, use_octf in itertools.product(
+            (torch.float, torch.double), (False, True)
+        ):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            octf = Standardize(m=1) if use_octf else None
+            train_X = torch.rand(5, 1, **tkwargs)
+            train_Y = torch.sin(train_X)
+            model = SimpleGPyTorchModel(train_X, train_Y, octf)
+
+            # must predict before conitioning on observations
+            model.posterior(torch.rand(2, 1, **tkwargs))
+
+            # test conditioning on observations
+            cm = model.condition_on_observations(
+                torch.rand(2, 1, **tkwargs), torch.rand(2, 1, **tkwargs)
+            )
+            self.assertIsInstance(cm, SimpleGPyTorchModel)
+            self.assertEqual(cm.train_targets.shape, torch.Size([7]))
+            model = SimpleGPyTorchModel(torch.rand(5, 1), torch.randn(5, 1))
+
+        model.posterior(torch.rand(2, 1))  # evaluate the model to form caches.
+        # Outside of fantasize, the inputs are validated.
+        with self.assertWarnsRegex(
+            BotorchTensorDimensionWarning, "Non-strict enforcement of"
+        ):
+            model.condition_on_observations(torch.randn(2, 1), torch.randn(5, 2, 1))
+        # Inside of fantasize, the inputs are not validated.
+        with fantasize(), warnings.catch_warnings(record=True) as ws:
+            warnings.filterwarnings("always", category=BotorchTensorDimensionWarning)
+            model.condition_on_observations(torch.randn(2, 1), torch.randn(5, 2, 1))
+        self.assertFalse(any(w.category is BotorchTensorDimensionWarning for w in ws))
+
+    def test_condition_on_observations_batched(self):
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            train_X = torch.rand(5, 1, **tkwargs)
+            train_Y = torch.cat([torch.sin(train_X), torch.cos(train_X)], dim=-1)
+            model = SimpleBatchedMultiOutputGPyTorchModel(train_X, train_Y)
+
+            # must predict before conitioning on observations
+            model.posterior(torch.rand(2, 1, **tkwargs))
+
+            # test conditioning on observations
+            cm = model.condition_on_observations(
+                torch.rand(2, 1, **tkwargs), torch.rand(2, 2, **tkwargs)
+            )
+            self.assertIsInstance(cm, SimpleBatchedMultiOutputGPyTorchModel)
+            self.assertEqual(cm.train_targets.shape, torch.Size([2, 7]))
+
     def test_posterior_transform(self):
         tkwargs = {"device": self.device, "dtype": torch.double}
         train_X = torch.rand(5, 2, **tkwargs)
@@ -564,11 +588,6 @@ def test_model_list_gpytorch_model(self):
             )
             self.assertIsInstance(posterior, GPyTorchPosterior)
             self.assertEqual(posterior.mean.shape, torch.Size([2, 1]))
-            # conditioning is not implemented (see ModelListGP for tests)
-            with self.assertRaises(NotImplementedError):
-                model.condition_on_observations(
-                    X=torch.rand(2, 1, **tkwargs), Y=torch.rand(2, 2, **tkwargs)
-                )
 
     def test_input_transform(self):
         # test that the input transforms are applied properly to individual models
@@ -651,3 +670,23 @@ def test_posterior_transform(self):
         post_tf = ScalarizedPosteriorTransform(torch.ones(2, **tkwargs))
         post = model.posterior(torch.rand(3, 1, **tkwargs), posterior_transform=post_tf)
         self.assertEqual(post.mean.shape, torch.Size([3, 1]))
+
+    def test_condition_on_observations_model_list(self):
+        torch.manual_seed(12345)
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            train_X1, train_X2 = (
+                torch.rand(5, 1, **tkwargs),
+                torch.rand(5, 1, **tkwargs),
+            )
+            train_Y1 = torch.sin(train_X1)
+            train_Y2 = torch.cos(train_X2)
+            m1 = SimpleGPyTorchModel(train_X1, train_Y1)
+            m2 = SimpleGPyTorchModel(train_X2, train_Y2)
+            model = SimpleModelListGPyTorchModel(m1, m2)
+
+            # conditioning is not implemented (see ModelListGP for tests)
+            with self.assertRaises(NotImplementedError):
+                model.condition_on_observations(
+                    X=torch.rand(2, 1, **tkwargs), Y=torch.rand(2, 2, **tkwargs)
+                )

test/models/test_model_list_gp_regression.py

@@ -170,39 +170,6 @@ def _base_test_ModelListGP(
         if gpytorch_posterior_expected:
             self.assertIsInstance(posterior.distribution, MultivariateNormal)
 
-        # test condition_on_observations
-        f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
-        f_y = torch.rand(2, 2, **tkwargs)
-        if fixed_noise:
-            noise = 0.1 + 0.1 * torch.rand_like(f_y)
-            cond_kwargs = {"noise": noise}
-        else:
-            cond_kwargs = {}
-        cm = model.condition_on_observations(f_x, f_y, **cond_kwargs)
-        self.assertIsInstance(cm, ModelListGP)
-
-        # test condition_on_observations batched
-        f_x = [torch.rand(3, 2, 1, **tkwargs) for _ in range(2)]
-        f_y = torch.rand(3, 2, 2, **tkwargs)
-        cm = model.condition_on_observations(f_x, f_y, **cond_kwargs)
-        self.assertIsInstance(cm, ModelListGP)
-
-        # test condition_on_observations batched (fast fantasies)
-        f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
-        f_y = torch.rand(3, 2, 2, **tkwargs)
-        cm = model.condition_on_observations(f_x, f_y, **cond_kwargs)
-        self.assertIsInstance(cm, ModelListGP)
-
-        # test condition_on_observations (incorrect input shape error)
-        with self.assertRaises(BotorchTensorDimensionError):
-            model.condition_on_observations(
-                f_x, torch.rand(3, 2, 3, **tkwargs), **cond_kwargs
-            )
-
-        # test X having wrong size
-        with self.assertRaises(BotorchTensorDimensionError):
-            model.condition_on_observations(f_x[:1], f_y)
-
         # test posterior transform
         X = torch.rand(3, 1, **tkwargs)
         weights = torch.tensor([1, 2], **tkwargs)
@@ -218,6 +185,48 @@ def _base_test_ModelListGP(
 
         return model
 
+    def test_condition_on_observations(self) -> None:
+        for dtype, outcome_transform in itertools.product(
+            (torch.float, torch.double), ("None", "Standardize", "Log", "Chained")
+        ):
+            with self.subTest(dtype=dtype, outcome_transform=outcome_transform):
+                tkwargs = {"device": self.device, "dtype": dtype}
+                model = _get_model(
+                    fixed_noise=False, outcome_transform=outcome_transform, **tkwargs
+                )
+                # need to predict before conditioning
+                test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
+                _ = model.posterior(test_x)
+
+                # test condition_on_observations
+                f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
+                f_y = torch.rand(2, 2, **tkwargs)
+                cond_kwargs = {}
+                cm = model.condition_on_observations(f_x, f_y, **cond_kwargs)
+                self.assertIsInstance(cm, ModelListGP)
+
+                # test condition_on_observations batched
+                f_x = [torch.rand(3, 2, 1, **tkwargs) for _ in range(2)]
+                f_y = torch.rand(3, 2, 2, **tkwargs)
+                cm = model.condition_on_observations(f_x, f_y, **cond_kwargs)
+                self.assertIsInstance(cm, ModelListGP)
+
+                # test condition_on_observations batched (fast fantasies)
+                f_x = [torch.rand(2, 1, **tkwargs) for _ in range(2)]
+                f_y = torch.rand(3, 2, 2, **tkwargs)
+                cm = model.condition_on_observations(f_x, f_y, **cond_kwargs)
+                self.assertIsInstance(cm, ModelListGP)
+
+                # test condition_on_observations (incorrect input shape error)
+                with self.assertRaises(BotorchTensorDimensionError):
+                    model.condition_on_observations(
+                        f_x, torch.rand(3, 2, 3, **tkwargs), **cond_kwargs
+                    )
+
+                # test X having wrong size
+                with self.assertRaises(BotorchTensorDimensionError):
+                    model.condition_on_observations(f_x[:1], f_y)
+
     def test_ModelListGP(self) -> None:
         for dtype, outcome_transform in itertools.product(
             (torch.float, torch.double), ("None", "Standardize", "Log", "Chained")