Add handling for failed simulations

autoemulate/experimental/calibration/history_matching.py

@@ -407,18 +407,15 @@ def simulate(self, x: TensorLike) -> tuple[TensorLike, TensorLike]:
         tuple[TensorLike, TensorLike]
             Tensors of succesfully simulated input parameters and predictions.
         """
-        y = self.simulator.forward_batch(x).to(self.device)
+        # if simulation fails, returned y and x have fewer rows than input x
+        y, x = self.simulator.forward_batch_skip_failures(x)
+        y = y.to(self.device)
+        x = x.to(self.device)
 
-        # Filter out runs that simulator failed to return predictions for
-        # TODO: this assumes that simulator returns None if it fails (see #438)
-        valid_indices = [i for i, res in enumerate(y) if res is not None]
-        valid_x = x[valid_indices]
-        valid_y = y[valid_indices]
+        self.train_y = torch.cat([self.train_y, y], dim=0)
+        self.train_x = torch.cat([self.train_x, x], dim=0)
 
-        self.train_y = torch.cat([self.train_y, valid_y], dim=0)
-        self.train_x = torch.cat([self.train_x, valid_x], dim=0)
-
-        return valid_x, valid_y
+        return x, y
 
     def run(
         self,

autoemulate/experimental/learners/base.py

@@ -144,6 +144,7 @@ def fit(self, *args):
             # If x is not, we skip the point (typically for Stream learners)
             self.logger.info("Appending new training data and refitting emulator.")
             y_true = self.simulator.forward(x)
+            assert isinstance(y_true, TensorLike)
             self.x_train = torch.cat([self.x_train, x])
             self.y_train = torch.cat([self.y_train, y_true])
             self.emulator.fit(self.x_train, self.y_train)

autoemulate/experimental/simulations/base.py

@@ -105,7 +105,7 @@ def sample_inputs(
         return lhd.sample(n_samples)
 
     @abstractmethod
-    def _forward(self, x: TensorLike) -> TensorLike:
+    def _forward(self, x: TensorLike) -> TensorLike | None:
         """
         Abstract method to perform the forward simulation.
 
@@ -116,13 +116,13 @@ def _forward(self, x: TensorLike) -> TensorLike:
 
         Returns
         -------
-        TensorLike
+        TensorLike | None
             Simulated output tensor. Shape = (1, self.out_dim).
             For example, if the simulator outputs two simulated variables,
-            then the shape would be (1, 2).
+            then the shape would be (1, 2). None if the simulation fails.
         """
 
-    def forward(self, x: TensorLike) -> TensorLike:
+    def forward(self, x: TensorLike) -> TensorLike | None:
         """
         Generate samples from input data using the simulator. Combines the
         abstract method `_forward` with some validation checks.
@@ -135,61 +135,108 @@ def forward(self, x: TensorLike) -> TensorLike:
         Returns
         -------
         TensorLike
-            Simulated output tensor.
+            Simulated output tensor. None if the simulation failed.
         """
-        y = self.check_matrix(self._forward(self.check_matrix(x)))
-        x, y = self.check_pair(x, y)
-        return y
+        y = self._forward(self.check_matrix(x))
+        if isinstance(y, TensorLike):
+            y = self.check_matrix(y)
+            x, y = self.check_pair(x, y)
+            return y
+        return None
 
-    def forward_batch(self, samples: TensorLike) -> TensorLike:
-        """
-        Run multiple simulations with different parameters.
+    def forward_batch(self, x: TensorLike) -> TensorLike:
+        """Run multiple simulations with different parameters.
+
+        For infallible simulators that always succeed.
+        If your simulator might fail, use `forward_batch_skip_failures()` instead.
 
         Parameters
         ----------
-        samples: TensorLike
+        x: TensorLike
             Tensor of input parameters to make predictions for.
 
-        Returns:
+        Returns
         -------
         TensorLike
             Tensor of simulation results of shape (n_batch, self.out_dim).
+
+        Raises
+        ------
+        RuntimeError
+            If the number of simulations does not match the input.
+            Use `forward_batch_skip_failures()` to handle failures.
+        """
+        results, x_valid = self.forward_batch_skip_failures(x)
+
+        # Raise an error if the number of simulations does not match the input
+        if x.shape[0] != x_valid.shape[0]:
+            msg = (
+                "Some simulations failed. Use forward_batch_skip_failures() to handle "
+                "failures."
+            )
+            raise RuntimeError(msg)
+
+        return results
+
+    def forward_batch_skip_failures(
+        self, x: TensorLike
+    ) -> tuple[TensorLike, TensorLike]:
+        """Run multiple simulations, skipping any that fail.
+
+        For simulators where for some inputs the simulation can fail.
+        Failed simulations are skipped, and only successful results are returned
+        along with their corresponding input parameters.
+
+        Parameters
+        ----------
+        x: TensorLike
+            Tensor of input parameters to make predictions for.
+
+        Returns
+        -------
+        tuple[TensorLike, TensorLike]
+            Tuple of (simulation_results, valid_input_parameters).
+            Only successful simulations are included.
         """
-        self.logger.info("Running batch simulation for %d samples", len(samples))
+        self.logger.info("Running batch simulation for %d samples", len(x))
 
         results = []
         successful = 0
+        valid_idx = []
 
         # Process each sample with progress tracking
         for i in tqdm(
-            range(len(samples)),
+            range(len(x)),
             desc="Running simulations",
             disable=not self.progress_bar,
-            total=len(samples),
+            total=len(x),
             unit="sample",
             unit_scale=True,
         ):
-            logger.debug("Running simulation for sample %d/%d", i + 1, len(samples))
-            result = self.forward(samples[i : i + 1])
+            logger.debug("Running simulation for sample %d/%d", i + 1, len(x))
+            result = self.forward(x[i : i + 1])
             if result is not None:
                 results.append(result)
                 successful += 1
-                logger.debug("Simulation %d/%d successful", i + 1, len(samples))
+                valid_idx.append(i)
+                logger.debug("Simulation %d/%d successful", i + 1, len(x))
             else:
                 logger.warning(
-                    "Simulation %d/%d failed. Result is None.", i + 1, len(samples)
+                    "Simulation %d/%d failed. Result is None.", i + 1, len(x)
                 )
 
         # Report results
         self.logger.info(
             "Successfully completed %d/%d simulations (%.1f%%)",
             successful,
-            len(samples),
-            (successful / len(samples) * 100 if len(samples) > 0 else 0.0),
+            len(x),
+            (successful / len(x) * 100 if len(x) > 0 else 0.0),
         )
 
         # stack results into a 2D array on first dim using torch
-        return torch.cat(results, dim=0)
+        results_tensor = torch.cat(results, dim=0)
+
+        return results_tensor, x[valid_idx]
 
     def get_parameter_idx(self, name: str) -> int:
         """

docs/experimental/tutorials/simulator/01_custom_simulations.ipynb

@@ -23,7 +23,8 @@
     "- `output_variables`: A parameter set at class initiation. A list of strings that are the names of the outputs that the simulator will return.\n",
     "- `_forward`: An abstract method that must be implemented by the user. This method will define a single forward pass of the simulation, taking in the input parameters and returning the output variables. *There are some important rules for this method:*\n",
     "  - The input to the method must be a tensor of shape `(1, n)` where `n` is the number of input parameters.\n",
-    "  - The output of a method must be a tensor of shape `(1, m)` where `m` is the number of output variables.\n"
+    "  - The output of th method must be a tensor of shape `(1, m)` where `m` is the number of output variables. \n",
+    "  - If the simulation fails, it must output `None`.\n"
    ]
   },
   {

tests/experimental/test_experimental_base_simulator.py

@@ -16,7 +16,7 @@ def __init__(
         # Call parent constructor
         super().__init__(parameters_range, output_names)
 
-    def _forward(self, x: TensorLike) -> TensorLike:
+    def _forward(self, x: TensorLike) -> TensorLike | None:
         """
         Implement abstract _forward method with a simple transformation.
         Input shape: (n_samples, n_features)
@@ -164,17 +164,12 @@ def test_abstract_class():
 def test_handle_simulation_failure():
     """Test handling of simulation failures in forward_batch"""
 
-    # This test demonstrates that the current implementation doesn't properly handle
-    # None returns (which would be a good enhancement)
     class ThresholdSimulator(MockSimulator):
-        def _forward(self, x: TensorLike) -> TensorLike:
+        def _forward(self, x: TensorLike) -> TensorLike | None:
             # Only process inputs where the first value is > 0.5
             if x[0, 0] > 0.5:
                 return super()._forward(x)
-            # For other cases, let's return a tensor of zeros instead of None
-            # since the type system doesn't allow None returns
-            # TODO (#438): update to handle failed simulations
-            return torch.zeros((x.shape[0], len(self._output_names)))
+            return None
 
     # Create simulator with float parameters
     params = {"param1": (0.0, 1.0), "param2": (0.0, 1.0), "param3": (0.0, 1.0)}
@@ -193,7 +188,9 @@ def _forward(self, x: TensorLike) -> TensorLike:
 
     # This should process all samples without errors
     # We're just verifying it doesn't crash
-    results = simulator.forward_batch(batch)
+    results, valid_x = simulator.forward_batch_skip_failures(batch)
+    assert isinstance(results, TensorLike)
 
     # Verify results shape
-    assert results.shape == (4, 1)
+    assert results.shape == (2, 1)
+    assert valid_x.shape == (2, 3)

tests/experimental/test_experimental_bayesian_calibration.py

@@ -7,6 +7,7 @@
     Projectile,
     ProjectileMultioutput,
 )
+from autoemulate.experimental.types import TensorLike
 
 
 @pytest.mark.parametrize(
@@ -20,6 +21,7 @@ def test_hmc_single_output(n_obs, n_chains, n_samples):
     sim = Projectile()
     x = sim.sample_inputs(100)
     y = sim.forward_batch(x)
+    assert isinstance(y, TensorLike)
     gp = GaussianProcessExact(x, y)
     gp.fit(x, y)
 
@@ -59,6 +61,7 @@ def test_hmc_multiple_output(n_obs, n_chains, n_samples):
     sim = ProjectileMultioutput()
     x = sim.sample_inputs(100)
     y = sim.forward_batch(x)
+    assert isinstance(y, TensorLike)
     gp = GaussianProcessExact(x, y)
     gp.fit(x, y)
 

tests/experimental/test_experimental_history_matching.py

@@ -107,6 +107,7 @@ def test_run(device):
     simulator = Epidemic()
     x = simulator.sample_inputs(10)
     y = simulator.forward_batch(x)
+    assert isinstance(y, TensorLike)
 
     # Run history matching
     gp = GaussianProcessExact(x, y, device=device)
@@ -145,6 +146,7 @@ def test_run_max_tries():
     simulator = Epidemic()
     x = simulator.sample_inputs(10)
     y = simulator.forward_batch(x)
+    assert isinstance(y, TensorLike)
 
     # Run history matching
     gp = GaussianProcessExact(x, y)

tests/experimental/test_learners.py

@@ -8,6 +8,7 @@
 from autoemulate.experimental.learners import stream
 from autoemulate.experimental.simulations.base import Simulator
 from autoemulate.experimental.simulations.projectile import ProjectileMultioutput
+from autoemulate.experimental.types import TensorLike
 from tqdm import tqdm
 
 
@@ -22,6 +23,7 @@ def learners(
 ) -> Iterable:
     x_train = simulator.sample_inputs(n_initial_samples)
     y_train = simulator.forward_batch(x_train)
+    assert isinstance(y_train, TensorLike)
     yield stream.Random(
         simulator=simulator,
         emulator=GaussianProcessExact(x_train, y_train),