Merge pull request #938 from alan-turing-institute/fix_test_set

Let user specify test set

Author: radka-j

autoemulate/core/compare.py

@@ -55,6 +55,7 @@ def __init__(
         self,
         x: InputLike,
         y: InputLike,
+        test_data: tuple[InputLike, InputLike] | None = None,
         models: list[type[Emulator] | str] | None = None,
         x_transforms_list: list[list[Transform | dict]] | None = None,
         y_transforms_list: list[list[Transform | dict]] | None = None,
@@ -81,6 +82,9 @@ def __init__(
             Input features.
         y: InputLike or None
             Target values (not needed if x is a Dataset).
+        test_data: tuple[InputLike, InputLike] | None
+            Optional test data as a tuple (x_test, y_test). If None, a random split
+            from the provided data is used. Defaults to None.
         models: list[type[Emulator]] | None
             List of emulator classes to compare. If None, all available emulators
             are used.
@@ -164,7 +168,17 @@ def __init__(
         self.models = updated_models
         if random_seed is not None:
             set_random_seed(seed=random_seed)
-        self.train_val, self.test = self._random_split(self._convert_to_dataset(x, y))
+
+        if test_data is None:
+            self.train_val, self.test = self._random_split(
+                self._convert_to_dataset(x, y)
+            )
+        else:
+            self.train_val = self._convert_to_dataset(x, y)
+            test_x, test_y = self._move_tensors_to_device(
+                *self._convert_to_tensors(*test_data)
+            )
+            self.test = self._convert_to_dataset(test_x, test_y)
 
         # Run the compare method with the provided models
         if not self.models:

tests/core/test_compare.py

@@ -10,6 +10,7 @@
 from autoemulate.emulators import DEFAULT_EMULATORS
 from autoemulate.emulators.base import Emulator
 from torch.distributions import Transform
+from torch.utils.data import TensorDataset
 
 
 @pytest.mark.parametrize("device", SUPPORTED_DEVICES)
@@ -462,3 +463,29 @@ def __call__(
         metric_names = [m.name for m in result.test_metrics]
         assert "custom_r2" in metric_names
         assert "rmse" in metric_names
+
+
+def test_ae_with_fixed_test_data(sample_data_for_ae_compare):
+    """Test AutoEmulate with a fixed test dataset."""
+    x, y = sample_data_for_ae_compare
+    models: list[str | type[Emulator]] = ["mlp", "RandomForest"]
+
+    # Create fixed test set
+    test_size = 25
+    x_test, y_test = x[:test_size], y[:test_size]
+    x_train, y_train = x[test_size:], y[test_size:]
+
+    ae = AutoEmulate(
+        x_train,
+        y_train,
+        models=models,
+        test_data=(x_test, y_test),
+        n_iter=2,
+        n_splits=2,
+        model_params={},  # Skip tuning for speed
+    )
+
+    assert isinstance(ae.test, TensorDataset)
+    assert ae.test.tensors == (x_test, y_test)
+    assert isinstance(ae.train_val, TensorDataset)
+    assert ae.train_val.tensors == (x_train, y_train)