Merge pull request #30 from rohanbabbar04/numpyro

Add support for numpyro models in SBC

Author: aloctavodia

docs/examples/gallery/sbc.md

@@ -105,3 +105,47 @@ plot_ecdf_pit(sbc.simulations)
 ```
 
 :::::
+
+:::::{tab-item} Numpyro
+:sync: numpyro
+
+We define a Numpyro Model, we use the centered eight schools model.
+
+```{jupyter-execute}
+import numpyro
+import numpyro.distributions as dist
+from jax import random
+from numpyro.infer import NUTS
+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 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)
+# We use the NUTS sampler
+nuts_kernel = NUTS(eight_schools_cauchy_prior)
+```
+
+Pass the model to the `SBC` class, set the number of simulations to 100, and run the simulations. For numpyro model, 
+we pass in the ``data_dir`` parameter.
+
+```{jupyter-execute}
+sbc = simuk.SBC(nuts_kernel,
+    sample_kwargs={"num_warmup": 50, "num_samples": 75},
+    num_simulations=100,
+    data_dir={"J": 8, "sigma": sigma, "y": y},
+)
+sbc.run_simulations()
+```
+
+To compare the prior and posterior distributions, we will plot the results.
+We expect a uniform distribution, the gray envelope corresponds to the 94% credible interval.
+
+```{jupyter-execute}
+plot_ecdf_pit(sbc.simulations, 
+              pc_kwargs={'col_wrap':4},
+              plot_kwargs={"xlabel":False}
+)
+```

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 dynamic = ["version"]
 description = "Simulation based calibration and generation of synthetic data."
 dependencies = [
-  "pymc>=5.20",
+  "arviz>=0.20.0",
   "arviz_base>=0.4.0",
   "tqdm"
 ]

requirements-dev.txt

@@ -8,3 +8,4 @@ pymc>=5.20.1
 bambi>=0.13.0
 arviz_base>=0.4.0
 ruff==0.9.1
+numpyro>=0.17.0

requirements-docs.txt

@@ -9,3 +9,4 @@ sphinx_tabs
 sphinx-design
 numpydoc
 jupyter-sphinx
+numpyro>=0.17.0

simuk/sbc.py

@@ -4,8 +4,20 @@
 from copy import copy
 from importlib.metadata import version
 
+try:
+    import pymc as pm
+except ImportError:
+    pass
+try:
+    import jax
+    from numpyro.handlers import seed, trace
+    from numpyro.infer import MCMC, Predictive
+    from numpyro.infer.mcmc import MCMCKernel
+except ImportError:
+    pass
+
 import numpy as np
-import pymc as pm
+from arviz import from_numpyro
 from arviz_base import extract, from_dict
 from tqdm import tqdm
 
@@ -35,8 +47,8 @@ class SBC:
 
     Parameters
     ----------
-    model : function
-        A PyMC or Bambi model. If a PyMC model the data needs to be defined as
+    model : pymc.Model, bambi.Model or numpyro.infer.mcmc.MCMCKernel
+        A PyMC, Bambi model or Numpyro MCMC kernel. If a PyMC model the data needs to be defined as
         mutable data.
     num_simulations : int
         How many simulations to run
@@ -45,6 +57,9 @@ class SBC:
     seed : int (optional)
         Random seed. This persists even if running the simulations is
         paused for whatever reason.
+    data_dir : dict
+        Keyword arguments passed to numpyro model, intended for use when providing
+        an MCMC Kernel model.
 
     Example
     -------
@@ -61,39 +76,63 @@ class SBC:
 
     """
 
-    def __init__(
-        self,
-        model,
-        num_simulations=1000,
-        sample_kwargs=None,
-        seed=None,
-    ):
-        if isinstance(model, pm.Model):
+    def __init__(self, model, num_simulations=1000, sample_kwargs=None, seed=None, data_dir=None):
+        if hasattr(model, "basic_RVs") and isinstance(model, pm.Model):
             self.engine = "pymc"
             self.model = model
-        else:
+        elif hasattr(model, "formula"):
             self.engine = "bambi"
             model.build()
             self.bambi_model = model
             self.model = model.backend.model
             self.formula = model.formula
             self.new_data = copy(model.data)
-
-        self.observed_vars = [obs_rvs.name for obs_rvs in self.model.observed_RVs]
+        elif isinstance(model, MCMCKernel):
+            self.engine = "numpyro"
+            self.numpyro_model = model
+            self.model = self.numpyro_model.model
+            self.run_simulations = self._run_simulations_numpyro
+            self.data_dir = data_dir
+        else:
+            raise ValueError(
+                "model should be one of pymc.Model, bambi.Model, or numpyro.infer.mcmc.MCMCKernel"
+            )
         self.num_simulations = num_simulations
-
-        self.var_names = [v.name for v in self.model.free_RVs]
-
         if sample_kwargs is None:
             sample_kwargs = {}
-        sample_kwargs.setdefault("progressbar", False)
-        sample_kwargs.setdefault("compute_convergence_checks", False)
+        if self.engine == "numpyro":
+            sample_kwargs.setdefault("num_warmup", 1000)
+            sample_kwargs.setdefault("num_samples", 1000)
+            sample_kwargs.setdefault("progress_bar", False)
+        else:
+            sample_kwargs.setdefault("progressbar", False)
+            sample_kwargs.setdefault("compute_convergence_checks", False)
         self.sample_kwargs = sample_kwargs
-
-        self.simulations = {name: [] for name in self.var_names}
-        self._simulations_complete = 0
         self.seed = seed
         self._seeds = self._get_seeds()
+        self._extract_variable_names()
+        self.simulations = {name: [] for name in self.var_names}
+        self._simulations_complete = 0
+
+    def _extract_variable_names(self):
+        """Extract observed and free variables from the model."""
+        if self.engine == "numpyro":
+            with trace() as tr:
+                with seed(rng_seed=int(self._seeds[0])):
+                    self.numpyro_model.model(**self.data_dir)
+            self.var_names = [
+                name
+                for name, site in tr.items()
+                if site["type"] == "sample" and not site.get("is_observed", False)
+            ]
+            self.observed_vars = [
+                name
+                for name, site in tr.items()
+                if site["type"] == "sample" and site.get("is_observed", False)
+            ]
+        else:
+            self.observed_vars = [obs.name for obs in self.model.observed_RVs]
+            self.var_names = [v.name for v in self.model.free_RVs]
 
     def _get_seeds(self):
         """Set the random seed, and generate seeds for all the simulations."""
@@ -110,6 +149,15 @@ def _get_prior_predictive_samples(self):
             prior = extract(idata, group="prior", keep_dataset=True)
         return prior, prior_pred
 
+    def _get_prior_predictive_samples_numpyro(self):
+        """Generate samples to use for the simulations using numpyro."""
+        predictive = Predictive(self.model, num_samples=self.num_simulations)
+        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}
+        return prior, prior_pred
+
     def _get_posterior_samples(self, prior_predictive_draw):
         """Generate posterior samples conditioned to a prior predictive sample."""
         new_model = pm.observe(self.model, prior_predictive_draw)
@@ -121,6 +169,14 @@ def _get_posterior_samples(self, prior_predictive_draw):
         posterior = extract(check, group="posterior", keep_dataset=True)
         return posterior
 
+    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}
+        mcmc.run(rng_seed, **free_vars_data, **prior_predictive_draw)
+        return from_numpyro(mcmc)["posterior"]
+
     def _convert_to_datatree(self):
         self.simulations = from_dict(
             {"prior_sbc": self.simulations},
@@ -171,3 +227,30 @@ def run_simulations(self):
             }
             self._convert_to_datatree()
             progress.close()
+
+    @quiet_logging("numpyro")
+    def _run_simulations_numpyro(self):
+        """Run all the simulations for Numpyro Model."""
+        prior, prior_pred = self._get_prior_predictive_samples_numpyro()
+        progress = tqdm(
+            initial=self._simulations_complete,
+            total=self.num_simulations,
+        )
+        try:
+            while self._simulations_complete < self.num_simulations:
+                idx = self._simulations_complete
+                prior_predictive_draw = {k: v[idx] for k, v in prior_pred.items()}
+                posterior = self._get_posterior_samples_numpyro(prior_predictive_draw)
+                for name in self.var_names:
+                    self.simulations[name].append(
+                        (posterior[name].sel(chain=0) < prior[name][idx]).sum(axis=0).values
+                    )
+                self._simulations_complete += 1
+                progress.update()
+        finally:
+            self.simulations = {
+                k: np.stack(v[: self._simulations_complete])[None, :]
+                for k, v in self.simulations.items()
+            }
+            self._convert_to_datatree()
+            progress.close()

simuk/tests/test_sbc.py

@@ -1,8 +1,11 @@
 import bambi as bmb
 import numpy as np
+import numpyro
+import numpyro.distributions as dist
 import pandas as pd
 import pymc as pm
 import pytest
+from numpyro.infer import NUTS
 
 import simuk
 
@@ -32,3 +35,24 @@ def test_sbc(model):
     )
     sbc.run_simulations()
     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 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)
+
+    sbc = simuk.SBC(
+        NUTS(eight_schools_cauchy_prior),
+        data_dir={"J": 8, "sigma": sigma, "y": y},
+        num_simulations=10,
+        sample_kwargs={"num_warmup": 50, "num_samples": 25},
+    )
+    sbc.run_simulations()
+    assert "prior_sbc" in sbc.simulations