refactor benchmarks into simulators

Author: LarsKue

bayesflow/__init__.py

@@ -1,6 +1,5 @@
 from . import (
     approximators,
-    benchmarks,
     adapters,
     datasets,
     diagnostics,

bayesflow/benchmarks/__init__.py

@@ -4,7 +4,12 @@
 from .make_simulator import make_simulator
 from .model_comparison_simulator import ModelComparisonSimulator
 from .simulator import Simulator
-from .two_moons import TwoMoons
+
+from .benchmark_simulators import (
+    LotkaVolterra,
+    SIR,
+    TwoMoons,
+)
 
 from ..utils._docs import _add_imports_to_all
 

bayesflow/benchmarks/simulators/two_moons.py

@@ -1,2 +1,3 @@
-from .two_moons import TwoMoons
 from .lotka_volterra import LotkaVolterra
+from .sir import SIR
+from .two_moons import TwoMoons

bayesflow/simulators/__init__.py

@@ -2,9 +2,13 @@
 
 from bayesflow.types import Shape
 from bayesflow.utils import batched_call, tree_stack
+from bayesflow.utils.decorators import allow_batch_size
 
+from ..simulator import Simulator
 
-class BenchmarkSimulator:
+
+class BenchmarkSimulator(Simulator):
+    @allow_batch_size
     def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
         """Runs simulated benchmark and returns `batch_size` parameter
         and observation batches