overriding simulator outputs with auto-batched inputs

Author: Kucharssim

bayesflow/simulators/sequential_simulator.py

@@ -10,7 +10,7 @@
 class SequentialSimulator(Simulator):
     """Combines multiple simulators into one, sequentially."""
 
-    def __init__(self, simulators: Sequence[Simulator], expand_outputs: bool = True):
+    def __init__(self, simulators: Sequence[Simulator], expand_outputs: bool = True, replace_inputs: bool = True):
         """
         Initialize a SequentialSimulator.
 
@@ -22,10 +22,13 @@ def __init__(self, simulators: Sequence[Simulator], expand_outputs: bool = True)
         expand_outputs : bool, optional
             If True, 1D output arrays are expanded with an additional dimension at the end.
             Default is True.
+        replace_inputs : bool, optional
+            If True, **kwargs are auto-batched and replace simulator outputs.
         """
 
         self.simulators = simulators
         self.expand_outputs = expand_outputs
+        self.replace_inputs = replace_inputs
 
     @allow_batch_size
     def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
@@ -53,6 +56,14 @@ def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
         for simulator in self.simulators:
             data |= simulator.sample(batch_shape, **(kwargs | data))
 
+            if self.replace_inputs:
+                common_keys = set(data.keys()) & set(kwargs.keys())
+                for key in common_keys:
+                    value = kwargs.pop(key)
+                    if isinstance(data[key], np.ndarray):
+                        value = np.broadcast_to(value, data[key].shape)
+                    data[key] = value
+
         if self.expand_outputs:
             data = {
                 key: np.expand_dims(value, axis=-1) if np.ndim(value) == 1 else value for key, value in data.items()