Inherit benchmark simulators

bayesflow/benchmarks/__init__.py

@@ -1 +1,12 @@
 from .benchmark import Benchmark
+from .bernoulli_glm_raw import BernoulliGLMRaw
+from .bernoulli_glm import BernoulliGLM
+from .gaussian_linear_uniform import GaussianLinearUniform
+from .gaussian_linear import GaussianLinear
+from .gaussian_mixture import GaussianMixture
+from .inverse_kinematics import InverseKinematics
+from .lotka_volterra import LotkaVolterra
+from .sir import SIR
+from .slcp_distractors import SLCPDistractors
+from .slcp import SLCP
+from .two_moons import TwoMoons

bayesflow/benchmarks/benchmark.py

@@ -1,39 +1,33 @@
-import importlib
-from functools import partial
-
+import numpy as np
+from bayesflow.types import Tensor
 from bayesflow.utils import batched_call
 
 
 class Benchmark:
-    def __init__(self, name: str, **kwargs):
-        """
-        Currently supported benchmarks:
-        - bernoulli_glm
-        - bernoulli_glm_raw
-        - gaussian_linear
-        - gaussian_linear_uniform
-        - gaussian_mixture
-        - inverse_kinematics
-        - lotka_volterra
-        - sir
-        - slcp
-        - slcp_distractors
-        - two_moons
-
-        TODO: Docs
+    def sample(self, batch_size: int) -> dict[str, Tensor]:
+        """Runs simulated benchmark and returns `batch_size` parameter
+        and observation batches
+
+        Parameters
+        ----------
+        batch_size: int
+            Number of parameter-observation batches to simulate
+
+        Returns
+        -------
+        dict[str, Tensor]: simulated parameters and observables
+            with shapes (`batch_size`, ...)
         """
 
-        self.name = name
-        self.module = self.get_module(name)
-        self.simulator = partial(getattr(self.module, "simulator"), **kwargs.pop("prior_kwargs", {}))
+        return batched_call(self, (batch_size,))
+
+    def prior(self):
+        raise NotImplementedError
 
-    def sample(self, batch_size: int):
-        return batched_call(self.simulator, (batch_size,))
+    def observation_model(self, params: np.ndarray):
+        raise NotImplementedError
 
-    @staticmethod
-    def get_module(name):
-        try:
-            benchmark_module = importlib.import_module(f"bayesflow.benchmarks.{name}")
-            return benchmark_module
-        except Exception as error:
-            raise ModuleNotFoundError(f"{name} is not a known benchmark") from error
+    def __call__(self):
+        prior_draws = self.prior()
+        observables = self.observation_model(prior_draws)
+        return dict(parameters=prior_draws, observables=observables)

bayesflow/benchmarks/bernoulli_glm.py

@@ -1,91 +1,84 @@
 import numpy as np
 from scipy.special import expit
-
-
-# Global covariance matrix computed once for efficiency
-F = np.zeros((9, 9))
-for i in range(9):
-    F[i, i] = 1 + np.sqrt(i / 9)
-    if i >= 1:
-        F[i, i - 1] = -2
-    if i >= 2:
-        F[i, i - 2] = 1
-Cov = np.linalg.inv(F.T @ F)
-
-
-def simulator():
-    """Non-configurable simulator running with default settings."""
-    prior_draws = prior()
-    observables = observation_model(prior_draws)
-    return dict(parameters=prior_draws, observables=observables)
-
-
-def prior(rng: np.random.Generator = None):
-    """Generates a random draw from the custom prior over the 10
-    Bernoulli GLM parameters (1 intercept and 9 weights). Uses a
-    global covariance matrix `Cov` for the multivariate Gaussian prior
-    over the model weights, which is pre-computed for efficiency.
-
-    Parameters
-    ----------
-    rng   : np.random.Generator or None, default: None
-        An optional random number generator to use.
-
-    Returns
-    -------
-    params : np.ndarray of shape (10,)
-        A single draw from the prior.
-    """
-
-    if rng is None:
-        rng = np.random.default_rng()
-    beta = rng.normal(0, 2)
-    f = rng.multivariate_normal(np.zeros(9), Cov)
-    return np.append(beta, f)
-
-
-def observation_model(params: np.ndarray, T: int = 100, scale_by_T: bool = True, rng: np.random.Generator = None):
-    """Simulates data from the custom Bernoulli GLM likelihood, see
-    https://arxiv.org/pdf/2101.04653.pdf, Task T.5
-
-    Important: `scale_sum` should be set to False if the simulator is used
-    with variable `T` during training, otherwise the information of `T` will
-    be lost.
-
-    Parameters
-    ----------
-    params     : np.ndarray of shape (10,)
-        The vector of model parameters (`params[0]` is intercept, `params[i], i > 0` are weights).
-    T          : int, optional, default: 100
-        The simulated duration of the task (eq. the number of Bernoulli draws).
-    scale_by_T : bool, optional, default: True
-        A flag indicating whether to scale the summayr statistics by T.
-    rng        : np.random.Generator or None, default: None
-        An optional random number generator to use.
-
-    Returns
-    -------
-    x : np.ndarray of shape (10,)
-        The vector of sufficient summary statistics of the data.
-    """
-
-    # Use default RNG, if None provided
-    if rng is None:
-        rng = np.random.default_rng()
-
-    # Unpack parameters
-    beta, f = params[0], params[1:]
-
-    # Generate design matrix
-    V = rng.normal(size=(9, T))
-
-    # Draw from Bernoulli GLM
-    z = rng.binomial(n=1, p=expit(V.T @ f + beta))
-
-    # Compute and return (scaled) sufficient summary statistics
-    x1 = np.sum(z)
-    x_rest = V @ z
-    x = np.append(x1, x_rest)
-    if scale_by_T:
-        x /= T
-    return x
+from .benchmark import Benchmark
+
+
+class BernoulliGLM(Benchmark):
+    def __init__(self, T: int = 100, scale_by_T: bool = True, rng: np.random.Generator = None):
+        """Bernoulli GLM simulated benchmark
+        See: https://arxiv.org/pdf/2101.04653.pdf, Task T.5
+
+        Important: `scale_sum` should be set to False if the simulator is used
+        with variable `T` during training, otherwise the information of `T` will
+        be lost.
+
+        Parameters
+        ----------
+        T: int, optional, default: 100
+            The simulated duration of the task (eq. the number of Bernoulli draws).
+        scale_by_T: bool, optional, default: True
+            A flag indicating whether to scale the summayr statistics by T.
+        rng: np.random.Generator or None, optional, default: None
+            An optional random number generator to use.
+        """
+
+        self.T = T
+        self.scale_by_T = scale_by_T
+        self.rng = rng
+        if self.rng is None:
+            self.rng = np.random.default_rng()
+
+        # Covariance matrix computed once for efficiency
+        F = np.zeros((9, 9))
+        i = np.arange(9)
+        F[i, i] = 1 + np.sqrt(i / 9)
+        F[i[1:], i[:-1]] = -2
+        F[i[2:], i[:-2]] = 1
+        self.cov = np.linalg.inv(F.T @ F)
+
+    def prior(self):
+        """Generates a random draw from the custom prior over the 10
+        Bernoulli GLM parameters (1 intercept and 9 weights). Uses a
+        global covariance matrix `Cov` for the multivariate Gaussian prior
+        over the model weights, which is pre-computed for efficiency.
+
+        Returns
+        -------
+        params: np.ndarray of shape (10, )
+            A single draw from the prior.
+        """
+
+        beta = self.rng.normal(0, 2)
+        f = self.rng.multivariate_normal(np.zeros(9), self.cov)
+        return np.append(beta, f)
+
+    def observation_model(self, params: np.ndarray):
+        """Simulates data from the custom Bernoulli GLM likelihood.
+
+        Parameters
+        ----------
+        params: np.ndarray of shape (10, )
+            The vector of model parameters (`params[0]` is intercept, `params[i], i > 0` are weights).
+
+        Returns
+        -------
+        x: np.ndarray of shape (10, )
+            The vector of sufficient summary statistics of the data.
+        """
+
+        # Unpack parameters
+        beta, f = params[0], params[1:]
+
+        # Generate design matrix
+        V = self.rng.normal(size=(9, self.T))
+
+        # Draw from Bernoulli GLM
+        z = self.rng.binomial(n=1, p=expit(V.T @ f + beta))
+
+        # Compute and return (scaled) sufficient summary statistics
+        x1 = np.sum(z)
+        x_rest = V @ z
+        x = np.append(x1, x_rest)
+        if self.scale_by_T:
+            x /= self.T
+        return x