Merge pull request #949 from cmsamaaa/feat/torch-simulator-device-support

Add a device-aware TorchSimulator base class

Author: radka-j

autoemulate/simulations/base.py

@@ -5,8 +5,9 @@
 from scipy.stats import qmc
 from tqdm import tqdm
 
+from autoemulate.core.device import TorchDeviceMixin
 from autoemulate.core.logging_config import get_configured_logger
-from autoemulate.core.types import TensorLike
+from autoemulate.core.types import DeviceLike, TensorLike
 from autoemulate.data.utils import ValidationMixin, set_random_seed
 
 logger = logging.getLogger("autoemulate")
@@ -362,3 +363,99 @@ def get_outputs_as_dict(self) -> dict[str, TensorLike]:
             output_dict[output_name] = self.results_tensor[:, i]
 
         return output_dict
+
+
+class TorchSimulator(Simulator, TorchDeviceMixin):
+    """
+    Simulator that runs computations on a specified torch device.
+
+    This subclass extends :class:`Simulator` with the :class:`TorchDeviceMixin`
+    so that simulators implemented in PyTorch (e.g., ``torchcor``) can run on
+    CPU or accelerator hardware. Inputs are moved to ``self.device`` before the
+    forward pass and the resulting tensors are kept on the same device.
+    """
+
+    def __init__(
+        self,
+        parameters_range: dict[str, tuple[float, float]],
+        output_names: list[str],
+        log_level: str = "progress_bar",
+        device: DeviceLike | None = None,
+    ):
+        Simulator.__init__(self, parameters_range, output_names, log_level)
+        TorchDeviceMixin.__init__(self, device=device)
+
+    def sample_inputs(
+        self, n_samples: int, random_seed: int | None = None, method: str = "lhs"
+    ) -> TensorLike:
+        """
+        Sample inputs and move them to the simulator's device.
+
+        Parameters
+        ----------
+        n_samples: int
+            Number of samples to generate.
+        random_seed: int | None
+            Optional random seed to make sampling reproducible.
+        method: str
+            Sampling method, one of ``"lhs"`` or ``"sobol"``.
+
+        Returns
+        -------
+        TensorLike
+            Sampled inputs located on ``self.device``.
+        """
+        samples = super().sample_inputs(
+            n_samples, random_seed=random_seed, method=method
+        )
+        (samples_device,) = self._move_tensors_to_device(samples)
+        return samples_device
+
+    def forward(self, x: TensorLike, allow_failures: bool = True) -> TensorLike | None:
+        """
+        Run a single simulation on the configured device.
+
+        Parameters
+        ----------
+        x: TensorLike
+            Input tensor with shape ``(n_samples, in_dim)``.
+        allow_failures: bool
+            When True, failures return ``None`` instead of raising.
+
+        Returns
+        -------
+        TensorLike | None
+            Simulation result on ``self.device`` or ``None`` on failure.
+        """
+        (x_device,) = self._move_tensors_to_device(x)
+        y = super().forward(x_device, allow_failures=allow_failures)
+        if isinstance(y, torch.Tensor):
+            return y.to(self.device)
+        return y
+
+    def forward_batch(
+        self, x: TensorLike, allow_failures: bool = True
+    ) -> tuple[TensorLike, TensorLike]:
+        """
+        Run a batch of simulations with device management.
+
+        Parameters
+        ----------
+        x: TensorLike
+            Batch of inputs with shape ``(batch_size, in_dim)``.
+        allow_failures: bool
+            Whether to skip failures (True) or raise immediately (False).
+
+        Returns
+        -------
+        tuple[TensorLike, TensorLike]
+            Tuple of ``(results, valid_inputs)`` residing on ``self.device``.
+        """
+        (x_device,) = self._move_tensors_to_device(x)
+        results, valid_inputs = super().forward_batch(
+            x_device, allow_failures=allow_failures
+        )
+        results = results.to(self.device)
+        self.results_tensor = results
+        valid_inputs = valid_inputs.to(self.device)
+        return results, valid_inputs

tests/simulations/test_base_simulator.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 from autoemulate.core.types import TensorLike
-from autoemulate.simulations.base import Simulator
+from autoemulate.simulations.base import Simulator, TorchSimulator
 from torch import Tensor
 
 
@@ -33,6 +33,25 @@ def _forward(self, x: TensorLike) -> TensorLike | None:
         return torch.cat(outputs, dim=1)
 
 
+class TorchMockSimulator(TorchSimulator):
+    """Torch-based simulator for testing TorchSimulator behaviour."""
+
+    def __init__(
+        self,
+        parameters_range: dict[str, tuple[float, float]],
+        output_names: list[str],
+        device: str | torch.device | None = "cpu",
+    ):
+        super().__init__(parameters_range, output_names, device=device)
+
+    def _forward(self, x: TensorLike) -> TensorLike | None:
+        outputs = []
+        for i, _ in enumerate(self._output_names):
+            output = torch.sum(x, dim=1) * (i + 1)
+            outputs.append(output.view(-1, 1))
+        return torch.cat(outputs, dim=1)
+
+
 @pytest.fixture
 def parameters_range():
     """Create test parameter ranges"""
@@ -324,3 +343,63 @@ def test_forward_batch_allow_failures():
     assert len(results) == 2  # All simulations successful
     assert len(valid_inputs) == 2
     assert torch.allclose(valid_inputs, success_batch)
+
+
+def test_torch_simulator_initializes_device(parameters_range):
+    sim = TorchMockSimulator(parameters_range, ["var1", "var2"], device="cpu")
+    assert sim.device == torch.device("cpu")
+
+
+def test_torch_simulator_forward_moves_inputs(parameters_range, monkeypatch):
+    sim = TorchMockSimulator(parameters_range, ["var1", "var2"], device="cpu")
+    original_move = TorchSimulator._move_tensors_to_device
+    calls = {"count": 0}
+
+    def recording_move(self, *args):
+        calls["count"] += 1
+        return original_move(self, *args)
+
+    monkeypatch.setattr(TorchSimulator, "_move_tensors_to_device", recording_move)
+    sim.forward(torch.tensor([[0.1, 0.2, 0.3]], dtype=torch.float32))
+    assert calls["count"] == 1
+
+
+def test_torch_simulator_forward_batch_moves_inputs(parameters_range, monkeypatch):
+    sim = TorchMockSimulator(parameters_range, ["var1"], device="cpu")
+    original_move = TorchSimulator._move_tensors_to_device
+    calls = {"count": 0}
+
+    def recording_move(self, *args):
+        calls["count"] += 1
+        return original_move(self, *args)
+
+    monkeypatch.setattr(TorchSimulator, "_move_tensors_to_device", recording_move)
+    batch = torch.tensor(
+        [
+            [0.2, 0.5, 0.5],
+            [0.6, 0.5, 0.5],
+        ],
+        dtype=torch.float32,
+    )
+    sim.forward_batch(batch)
+    # Forward batch moves the entire tensor once, then each per-sample forward
+    # moves the corresponding slice.
+    assert calls["count"] == 1 + len(batch)
+
+
+def test_torch_simulator_sample_inputs_on_device(parameters_range, monkeypatch):
+    sim = TorchMockSimulator(parameters_range, ["var1"], device="cpu")
+    original_move = TorchSimulator._move_tensors_to_device
+    calls = {"count": 0}
+
+    def recording_move(self, *args):
+        calls["count"] += 1
+        return original_move(self, *args)
+
+    monkeypatch.setattr(
+        TorchSimulator,
+        "_move_tensors_to_device",
+        recording_move,
+    )
+    sim.sample_inputs(5)
+    assert calls["count"] == 1