Merge branch 'alan-turing-institute:main' into plot_loss

Author: TheCSGuy25

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:

autoemulate/emulators/base.py

@@ -1,6 +1,5 @@
 import random
 from abc import ABC, abstractmethod
-from typing import ClassVar
 
 import numpy as np
 import torch
@@ -545,7 +544,6 @@ class PyTorchBackend(nn.Module, Emulator):
     batch_size: int = 16
     shuffle: bool = True
     epochs: int = 10
-    loss_history: ClassVar[list[float]] = []
     verbose: bool = False
     loss_fn: nn.Module = nn.MSELoss()
     optimizer_cls: type[optim.Optimizer] = optim.Adam
@@ -570,6 +568,8 @@ def _fit(self, x: TensorLike, y: TensorLike):
         y: OutputLike or None
             Target values (not needed if x is a DataLoader).
         """
+        self.loss_history: list[float] = []
+
         self.train()  # Set model to training mode
 
         # Convert input to DataLoader if not already

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)