Ax side fixes for adding feature & outcome names to datasets (#1853)

Summary:
Pull Request resolved: #1853

Follows up pytorch/botorch#2015 to propagate the changes in Ax and fixes the tests.

Reviewed By: Balandat

Differential Revision: D49340640

fbshipit-source-id: c775aaafaf35c84d8d342d1fcfae55f59e974f20

Author: saitcakmak

ax/benchmark/tests/test_benchmark.py

@@ -6,7 +6,6 @@
 import tempfile
 
 import numpy as np
-import torch
 from ax.benchmark.benchmark import (
     benchmark_multiple_problems_methods,
     benchmark_one_method_problem,
@@ -40,11 +39,10 @@
     get_sobol_benchmark_method,
     get_soo_surrogate,
 )
-from ax.utils.testing.core_stubs import get_experiment
+from ax.utils.testing.core_stubs import get_dataset, get_experiment
 from ax.utils.testing.mock import fast_botorch_optimize
 from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy
 from botorch.test_functions.synthetic import Branin
-from botorch.utils.datasets import SupervisedDataset
 
 
 class TestBenchmark(TestCase):
@@ -122,9 +120,7 @@ def test_replication_sobol_surrogate(self) -> None:
 
             with self.subTest(name, problem=problem):
                 surrogate, datasets = not_none(problem.get_surrogate_and_datasets)()
-                datasets = [
-                    SupervisedDataset(X=torch.zeros((2, 2)), Y=torch.zeros((2, 2)))
-                ]
+                datasets = [get_dataset()]
                 surrogate.fit(
                     datasets,
                     metric_names=[],

ax/modelbridge/pairwise.py

@@ -63,7 +63,12 @@ def _convert_observations(
             # pyre-fixme[6]: For 2nd param expected `LongTensor` but got `Tensor`.
             dataset_X = SliceContainer(datapoints, comparisons, event_shape=event_shape)
             dataset_Y = torch.tensor([[0, 1]]).expand(comparisons.shape)
-            dataset = RankingDataset(X=dataset_X, Y=dataset_Y)
+            dataset = RankingDataset(
+                X=dataset_X,
+                Y=dataset_Y,
+                feature_names=parameters,
+                outcome_names=[outcome],
+            )
 
             datasets.append(dataset)
             candidate_metadata.append(candidate_metadata_dict[outcome])

ax/modelbridge/tests/test_torch_modelbridge.py

@@ -82,11 +82,15 @@ def test_TorchModelBridge(self, mock_init, dtype=None, device=None) -> None:
                 X=X,
                 Y=torch.tensor([[3.0], [1.0]], **tkwargs),
                 Yvar=torch.tensor([[4.0], [2.0]], **tkwargs),
+                feature_names=feature_names,
+                outcome_names=["y1"],
             ),
             "y2": SupervisedDataset(
                 X=X,
                 Y=torch.tensor([[2.0], [0.0]], **tkwargs),
                 Yvar=torch.tensor([[2.0], [1.0]], **tkwargs),
+                feature_names=feature_names,
+                outcome_names=["y2"],
             ),
         }
         observation_features = [

ax/modelbridge/torch.py

@@ -329,9 +329,16 @@ def _convert_observations(
                 Yvars[outcome], dtype=self.dtype, device=self.device
             ).unsqueeze(-1)
             if Yvar.isnan().all():
-                dataset = SupervisedDataset(X=X, Y=Y)
+                Yvar = None
             else:
-                dataset = SupervisedDataset(X=X, Y=Y, Yvar=Yvar.clamp_min(1e-6))
+                Yvar = Yvar.clamp_min(1e-6)
+            dataset = SupervisedDataset(
+                X=X,
+                Y=Y,
+                Yvar=Yvar,
+                feature_names=parameters,
+                outcome_names=[outcome],
+            )
             datasets.append(dataset)
             candidate_metadata.append(candidate_metadata_dict[outcome])
 

ax/models/tests/test_alebo.py

@@ -266,7 +266,7 @@ def testAcq(self) -> None:
         self.assertTrue(torch.allclose(Z, X[:, 0, :]))
 
     @fast_botorch_optimize
-    def testALEBO(self) -> None:
+    def test_ALEBO(self) -> None:
         B = torch.tensor(
             [[1.0, 2.0, 3.0, 4.0, 5.0], [2.0, 3.0, 4.0, 5.0, 6.0]], dtype=torch.double
         )
@@ -288,7 +288,13 @@ def testALEBO(self) -> None:
         )
         train_Y = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.double)
         train_Yvar = 0.1 * torch.ones(3, 1, dtype=torch.double)
-        dataset = SupervisedDataset(X=train_X, Y=train_Y, Yvar=train_Yvar)
+        dataset = SupervisedDataset(
+            X=train_X,
+            Y=train_Y,
+            Yvar=train_Yvar,
+            feature_names=[f"x{i}" for i in range(5)],
+            outcome_names=["y"],
+        )
 
         # Test fit
         m.fit(
@@ -368,7 +374,13 @@ def testALEBO(self) -> None:
             ],
             dtype=torch.double,
         )
-        dataset2 = SupervisedDataset(X=train_X2, Y=train_Y, Yvar=train_Yvar)
+        dataset2 = SupervisedDataset(
+            X=train_X2,
+            Y=train_Y,
+            Yvar=train_Yvar,
+            feature_names=[f"x{i}" for i in range(5)],
+            outcome_names=["y2"],
+        )
         m.update(datasets=[dataset, dataset2])
         self.assertTrue(torch.allclose(m.Xs[0], (B @ train_X.t()).t()))
         self.assertTrue(torch.allclose(m.Xs[1], (B @ train_X2.t()).t()))

ax/models/tests/test_botorch_kg.py

@@ -38,14 +38,16 @@ def setUp(self) -> None:
             "device": torch.device("cpu"),
             "dtype": torch.double,
         }
+        self.feature_names = ["x1", "x2", "x3"]
+        self.metric_names = ["y"]
         self.dataset = SupervisedDataset(
             X=torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **self.tkwargs),
             Y=torch.tensor([[3.0], [4.0]], **self.tkwargs),
             Yvar=torch.tensor([[0.0], [2.0]], **self.tkwargs),
+            feature_names=self.feature_names,
+            outcome_names=self.metric_names,
         )
         self.bounds = [(0.0, 1.0), (1.0, 4.0), (2.0, 5.0)]
-        self.feature_names = ["x1", "x2", "x3"]
-        self.metric_names = ["y"]
         self.acq_options = {"num_fantasies": 30, "mc_samples": 30}
         self.objective_weights = torch.tensor([1.0], **self.tkwargs)
         self.optimizer_options = {

ax/models/tests/test_botorch_mes.py

@@ -30,16 +30,18 @@ def setUp(self) -> None:
             "device": torch.device("cpu"),
             "dtype": torch.double,
         }
+        self.feature_names = ["x1", "x2", "x3"]
+        self.metric_names = ["y"]
         self.training_data = [
             SupervisedDataset(
                 X=torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **self.tkwargs),
                 Y=torch.tensor([[3.0], [4.0]], **self.tkwargs),
                 Yvar=torch.tensor([[0.0], [2.0]], **self.tkwargs),
+                feature_names=self.feature_names,
+                outcome_names=self.metric_names,
             )
         ]
         self.bounds = [(0.0, 1.0), (1.0, 4.0), (2.0, 5.0)]
-        self.feature_names = ["x1", "x2", "x3"]
-        self.metric_names = ["y"]
         self.acq_options = {"num_fantasies": 30, "candidate_size": 100}
         self.objective_weights = torch.tensor([1.0], **self.tkwargs)
         self.optimize_acqf = "ax.models.torch.botorch_mes.optimize_acqf"