Add simulator function to SBC

simuk/sbc.py

@@ -17,6 +17,7 @@
     pass
 
 import numpy as np
+import xarray as xr
 from arviz_base import extract, from_dict, from_numpyro
 from tqdm import tqdm
 
@@ -59,6 +60,9 @@ class SBC:
     data_dir : dict
         Keyword arguments passed to numpyro model, intended for use when providing
         an MCMC Kernel model.
+        simulator : callable
+            A custom simulator function that takes as input the model parameters and
+            a int parameter named `seed`, and returns a dictionary of named observations.
 
     Example
     -------
@@ -73,7 +77,15 @@ class SBC:
         sbc.run_simulations()
     """
 
-    def __init__(self, model, num_simulations=1000, sample_kwargs=None, seed=None, data_dir=None):
+    def __init__(
+        self,
+        model,
+        num_simulations=1000,
+        sample_kwargs=None,
+        seed=None,
+        data_dir=None,
+        simulator=None,
+    ):
         if hasattr(model, "basic_RVs") and isinstance(model, pm.Model):
             self.engine = "pymc"
             self.model = model
@@ -110,6 +122,22 @@ def __init__(self, model, num_simulations=1000, sample_kwargs=None, seed=None, d
         self._extract_variable_names()
         self.simulations = {name: [] for name in self.var_names}
         self._simulations_complete = 0
+        if simulator is not None and not callable(simulator):
+            raise ValueError("simulator should be a function or None")
+        if simulator is not None and self.observed_vars:
+            logging.warning(
+                "Provided model contains both observed variables and a simulator. "
+                "Ignoring observed variables and using the simulator instead."
+            )
+        if simulator is None and not self.observed_vars and self.engine == "pymc":
+            # Ideally, we could raise an error early for `numpyro` also,
+            # but `factor` also produces 'observed_vars'
+            raise ValueError(
+                "There are no observed variables, and PyMC will not generate prior "
+                "predictive samples. Either change the model or specify a simulator "
+                "with the `simulator` argument."
+            )
+        self.simulator = simulator
 
     def _extract_variable_names(self):
         """Extract observed and free variables from the model."""
@@ -142,8 +170,42 @@ def _get_prior_predictive_samples(self):
             idata = pm.sample_prior_predictive(
                 samples=self.num_simulations, random_seed=self._seeds[0]
             )
-            prior_pred = extract(idata, group="prior_predictive", keep_dataset=True)
             prior = extract(idata, group="prior", keep_dataset=True)
+            if self.simulator is None:
+                prior_pred = extract(idata, group="prior_predictive", keep_dataset=True)
+                return prior, prior_pred
+            # Deal with custom simulator
+            prior_pred = []
+            for i in range(prior.sizes["sample"]):
+                params = {var: prior[var].isel(sample=i).values for var in prior.data_vars}
+                params["seed"] = self._seeds[i]
+                try:
+                    res = self.simulator(**params)
+                    assert isinstance(
+                        res, dict
+                    ), f"Simulator must return a dictionary, got {type(res)}"
+                    prior_pred.append(res)
+                except Exception as e:
+                    raise ValueError(
+                        f"Error generating prior predictive sample with parameters {params}: {e}."
+                    )
+            # --- Convert list of dicts to xarray.Dataset ---
+            # Get the sample coordinate (keeps chain/draw MultiIndex intact)
+            sample_coord = prior.coords["sample"]
+
+            # Collect variables into dict-of-arrays, stacking along 'sample'
+            data_vars = {}
+            for key in prior_pred[0].keys():
+                stacked = np.stack([np.asarray(d[key]) for d in prior_pred], axis=-1)
+                dims = [f"{key}_dim_{i}" for i in range(stacked.ndim - 1)] + ["sample"]
+                data_vars[key] = (dims, stacked)
+
+            # Build dataset
+            prior_pred = xr.Dataset(
+                data_vars=data_vars,
+                coords={**{k: v for k, v in prior.coords.items()}, "sample": sample_coord},
+                attrs=prior.attrs,
+            )
         return prior, prior_pred
 
     def _get_prior_predictive_samples_numpyro(self):
@@ -152,7 +214,15 @@ def _get_prior_predictive_samples_numpyro(self):
         free_vars_data = {k: v for k, v in self.data_dir.items() if k not in self.observed_vars}
         samples = predictive(jax.random.PRNGKey(self._seeds[0]), **free_vars_data)
         prior = {k: v for k, v in samples.items() if k not in self.observed_vars}
-        prior_pred = {k: v for k, v in samples.items() if k in self.observed_vars}
+        if self.simulator:
+            results = []
+            for i, vals in enumerate(zip(*prior.values())):
+                params = dict(zip(prior.keys(), vals))
+                params["seed"] = self._seeds[i]
+                results.append(self.simulator(**params))
+            prior_pred = dict(zip(results[0].keys(), zip(*[result.values() for result in results])))
+        else:
+            prior_pred = {k: v for k, v in samples.items() if k in self.observed_vars}
         return prior, prior_pred
 
     def _get_posterior_samples(self, prior_predictive_draw):
@@ -170,7 +240,12 @@ def _get_posterior_samples_numpyro(self, prior_predictive_draw):
         """Generate posterior samples using numpyro conditioned to a prior predictive sample."""
         mcmc = MCMC(self.numpyro_model, **self.sample_kwargs)
         rng_seed = jax.random.PRNGKey(self._seeds[self._simulations_complete])
-        free_vars_data = {k: v for k, v in self.data_dir.items() if k not in self.observed_vars}
+        # If using a custom simulator, some variables present in `prior_predictive_draw`
+        # might be missing from self.observed_vars.
+        # TODO: Not sure if the union is redundant here and perhaps prior_predictive_draw.keys()
+        # could be sufficient.
+        extended_observed_vars = set(prior_predictive_draw.keys()).union(self.observed_vars)
+        free_vars_data = {k: v for k, v in self.data_dir.items() if k not in extended_observed_vars}
         mcmc.run(rng_seed, **free_vars_data, **prior_predictive_draw)
         return from_numpyro(mcmc)["posterior"]
 

simuk/tests/test_sbc.py

@@ -1,4 +1,5 @@
 import bambi as bmb
+import jax.numpy as jnp
 import numpy as np
 import numpyro
 import numpyro.distributions as dist
@@ -11,23 +12,66 @@
 
 np.random.seed(1234)
 
+# Test data
 data = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
 sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
 
+# PyMC models
 with pm.Model() as centered_eight:
     mu = pm.Normal("mu", mu=0, sigma=5)
     tau = pm.HalfCauchy("tau", beta=5)
     theta = pm.Normal("theta", mu=mu, sigma=tau, shape=8)
     y_obs = pm.Normal("y", mu=theta, sigma=sigma, observed=data)
 
+with pm.Model() as centered_eight_no_observed:
+    mu = pm.Normal("mu", mu=0, sigma=5)
+    tau = pm.HalfCauchy("tau", beta=5)
+    theta = pm.Normal("theta", mu=mu, sigma=tau, shape=8)
+
+    def log_likelihood(theta, observed):
+        return pm.math.sum(pm.logp(pm.Normal.dist(mu=theta, sigma=sigma), observed))
+
+    pm.Potential("y_loglike", log_likelihood(mu, data))
+
+# Bambi model
 x = np.random.normal(0, 1, 20)
 y = 2 + np.random.normal(x, 1)
 df = pd.DataFrame({"x": x, "y": y})
 bmb_model = bmb.Model("y ~ x", df)
 
+# NumPyro models
+def eight_schools_cauchy_prior(J, sigma, y=None):
+    mu = numpyro.sample("mu", dist.Normal(0, 5))
+    tau = numpyro.sample("tau", dist.HalfCauchy(5))
+    with numpyro.plate("J", J):
+        theta = numpyro.sample("theta", dist.Normal(mu, tau))
+    numpyro.sample("y", dist.Normal(theta, sigma), obs=y)
+
+
+def eight_schools_cauchy_prior_no_observed(J, sigma, y=None):
+    mu = numpyro.sample("mu", dist.Normal(0, 5))
+    tau = numpyro.sample("tau", dist.HalfCauchy(5))
+    with numpyro.plate("J", J):
+        theta = numpyro.sample("theta", dist.Normal(mu, tau))
+    if y is not None:
+        log_likelihood = jnp.sum(dist.Normal(theta, sigma).log_prob(y))
+        numpyro.factor("custom_likelihood", log_likelihood)
 
+
+# Custom simulator functions
+def centered_eight_simulator(theta, seed, **kwargs):
+    rng = np.random.default_rng(seed)
+    return {"y": rng.normal(theta, sigma)}
+
+
+def bmb_simulator(mu, sigma, seed, **kwargs):
+    rng = np.random.default_rng(seed)
+    return {"y": rng.normal(mu, sigma)}
+
+
+# --- Tests with observed variables ---
 @pytest.mark.parametrize("model", [centered_eight, bmb_model])
-def test_sbc(model):
+def test_sbc_with_observed_data(model):
     sbc = simuk.SBC(
         model,
         num_simulations=10,
@@ -37,22 +81,76 @@ def test_sbc(model):
     assert "prior_sbc" in sbc.simulations
 
 
-def test_sbc_numpyro():
-    y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
-    sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
+def test_sbc_numpyro_with_observed_data():
+    sbc = simuk.SBC(
+        NUTS(eight_schools_cauchy_prior),
+        data_dir={"J": 8, "sigma": sigma, "y": data},
+        num_simulations=10,
+        sample_kwargs={"num_warmup": 50, "num_samples": 25},
+    )
+    sbc.run_simulations()
+    assert "prior_sbc" in sbc.simulations
 
-    def eight_schools_cauchy_prior(J, sigma, y=None):
-        mu = numpyro.sample("mu", dist.Normal(0, 5))
-        tau = numpyro.sample("tau", dist.HalfCauchy(5))
-        with numpyro.plate("J", J):
-            theta = numpyro.sample("theta", dist.Normal(mu, tau))
-        numpyro.sample("y", dist.Normal(theta, sigma), obs=y)
 
+# --- Tests with custom simulators ---
+@pytest.mark.parametrize(
+    "model,simulator",
+    [
+        # Case 1: Both simulator function and observed variables present
+        (centered_eight, centered_eight_simulator),
+        # Case 2: Only simulator function present
+        (centered_eight_no_observed, centered_eight_simulator),
+        # Case 3: bambi model with custom simulator
+        (bmb_model, bmb_simulator),
+    ],
+)
+def test_sbc_with_custom_simulator(model, simulator):
     sbc = simuk.SBC(
-        NUTS(eight_schools_cauchy_prior),
-        data_dir={"J": 8, "sigma": sigma, "y": y},
+        model, num_simulations=10, sample_kwargs={"draws": 5, "tune": 5}, simulator=simulator
+    )
+    sbc.run_simulations()
+    assert "prior_sbc" in sbc.simulations
+
+
+@pytest.mark.parametrize(
+    "model,simulator",
+    [
+        # Case 1: Both simulator function and observed variables present
+        (eight_schools_cauchy_prior, centered_eight_simulator),
+        # Case 2: Only simulator function present
+        (eight_schools_cauchy_prior_no_observed, centered_eight_simulator),
+    ],
+)
+def test_sbc_numpyro_with_custom_simulator(model, simulator):
+    sbc = simuk.SBC(
+        NUTS(model),
+        data_dir={"J": 8, "sigma": sigma, "y": data},
         num_simulations=10,
         sample_kwargs={"num_warmup": 50, "num_samples": 25},
+        simulator=simulator,
     )
     sbc.run_simulations()
     assert "prior_sbc" in sbc.simulations
+
+
+# --- Error handling tests with custom simulators ---
+def test_sbc_fail_no_observed_variable():
+    with pytest.raises(ValueError, match="no observed variables"):
+        simuk.SBC(
+            centered_eight_no_observed,
+            num_simulations=10,
+            sample_kwargs={"draws": 5, "tune": 5},
+        )
+
+
+def test_sbc_numpyro_fail_no_observed_variable():
+    # Note: factor variables are catalogued as 'observed_vars' in NumPyro
+    # therefore, we cannot raise an early exception with an informative message
+    with pytest.raises(ValueError):
+        sbc = simuk.SBC(
+            NUTS(eight_schools_cauchy_prior_no_observed),
+            data_dir={"J": 8, "sigma": sigma, "y": data},
+            num_simulations=10,
+            sample_kwargs={"num_warmup": 50, "num_samples": 25},
+        )
+        sbc.run_simulations()