Implement ContextualDataset.__eq__ (#3005)

Summary:
Pull Request resolved: #3005

--

Reviewed By: esantorella

Differential Revision: D81800547

fbshipit-source-id: cb2a50cc13e52cd4d0c3dad262bf6f171f1f29f7

Author: saitcakmak

botorch/utils/datasets.py

@@ -601,8 +601,6 @@ def __init__(
         self.datasets: dict[str, SupervisedDataset] = {
             ds.outcome_names[0]: ds for ds in datasets
         }
-        self.feature_names = datasets[0].feature_names
-        self.outcome_names = list(self.datasets.keys())
         self.parameter_decomposition = parameter_decomposition
         self.metric_decomposition = metric_decomposition
         self._validate_datasets()
@@ -614,6 +612,14 @@ def __init__(
         }
         self.group_indices = None
 
+    @property
+    def feature_names(self) -> list[str]:
+        return self.datasets[self.outcome_names[0]].feature_names
+
+    @property
+    def outcome_names(self) -> list[str]:
+        return list(self.datasets.keys())
+
     @property
     def X(self) -> Tensor:
         return self.datasets[self.outcome_names[0]].X
@@ -737,6 +743,14 @@ def _validate_decompositions(self) -> None:
                         f"{outcome} is missing in metric_decomposition."
                     )
 
+    def __eq__(self, other: Any) -> bool:
+        return (
+            type(other) is type(self)
+            and self.datasets == other.datasets
+            and self.parameter_decomposition == other.parameter_decomposition
+            and self.metric_decomposition == other.metric_decomposition
+        )
+
     def clone(
         self, deepcopy: bool = False, mask: Tensor | None = None
     ) -> ContextualDataset:

test/utils/test_datasets.py

@@ -729,3 +729,17 @@ def test_clone_contextual_dataset(self):
                     )
                 else:
                     self.assertIsNone(context_dt2.metric_decomposition)
+
+    def test_contextual_dataset_equality(self) -> None:
+        context_dt, _ = make_contextual_dataset(has_yvar=True, contextual_outcome=True)
+        clone = context_dt.clone()
+        self.assertEqual(context_dt, clone)
+        for yvar, outcome in (
+            (True, False),
+            (False, True),
+            (False, False),
+        ):
+            new_dt, _ = make_contextual_dataset(
+                has_yvar=yvar, contextual_outcome=outcome
+            )
+            self.assertNotEqual(context_dt, new_dt)