Posterior SBC Support

README.md

@@ -1,9 +1,13 @@
 # Simuk
 
 Simuk is a Python library for simulation-based calibration (SBC) and the generation of synthetic data.
-Simulation-Based Calibration is a method for validating Bayesian inference by checking whether the posterior distributions align with the expected theoretical results derived from the prior.
 
-Simuk works with [PyMC](http://docs.pymc.io), [Bambi](https://bambinos.github.io/bambi/) and [NumPyro](https://num.pyro.ai/en/latest/index.html) models.
+Prior Simulation-Based Calibration (Prior SBC) is a method for validating Bayesian inference by checking whether the posterior distributions align with the expected theoretical results derived from the prior.
+
+Posterior Simulation-Based Calibration (Posterior SBC) is a method for validating Bayesian inference by checking whether the posterior distributions conditioned on the augmented data (original + posterior predictive) align with the expected theoretical results derived from the posterior.
+
+For Prior SBC, Simuk works with [PyMC](http://docs.pymc.io), [Bambi](https://bambinos.github.io/bambi/) and [NumPyro](https://num.pyro.ai/en/latest/index.html) models.
+For Posterior SBC, Simuk only works with [PyMC](http://docs.pymc.io) models for now.
 
 ## Installation
 
@@ -15,6 +19,8 @@ pip install simuk
 
 ## Quickstart
 
+### Prior SBC
+
 1. Define a PyMC or Bambi model. For example, the centered eight schools model:
 
     ```python
@@ -35,7 +41,7 @@ pip install simuk
     ```python
     sbc = SBC(centered_eight,
             num_simulations=100, # ideally this should be higher, like 1000
-            sample_kwargs={'draws': 25, 'tune': 50})
+            sample_kwargs={'draws': 100, 'tune': 100})
 
     sbc.run_simulations()
     ```
@@ -52,7 +58,74 @@ should be close to uniform and within the oval envelope.
     );
     ```
 
-![Simulation based calibration plots, ecdf](ecdf.png)
+![Prior Simulation based calibration plots, ecdf](docs/examples/img/prior_sbc.png)
+
+We see that due to the funnel neck in the eight schools model, the inference algorithm is not well-calibrated, as indicated by the red points.
+
+### Posterior SBC
+
+Posterior SBC evaluates validity locally, conditional on observed data. It is
+currently implemented for PyMC. This requires storing observed data in
+`pm.Data` containers, using `dims` instead of static shapes, and resizing
+covariates and coords in an `update_data` callback to match the augmented data.
+
+1. Define the model with `pm.Data` and `dims`:
+
+    ```python
+    import numpy as np
+    import pymc as pm
+
+    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])
+
+    with pm.Model(coords={"school": np.arange(8)}) as centered_eight:
+        school_idx = pm.Data("school_idx", np.arange(8))
+        y_data = pm.Data("y_data", data)
+        sigma_data = pm.Data("sigma_data", sigma)
+
+        mu = pm.Normal("mu", mu=0, sigma=5)
+        tau = pm.HalfCauchy("tau", beta=5)
+        theta = pm.Normal("theta", mu=mu, sigma=tau, dims="school")
+        y_obs = pm.Normal("y", mu=theta[school_idx], sigma=sigma_data, observed=y_data)
+    ```
+
+2. Sample once to obtain the original trace:
+
+    ```python
+    with centered_eight:
+        idata = pm.sample(progressbar=False)
+    ```
+
+3. Define `update_data` to resize covariates and run Posterior SBC:
+
+    ```python
+    import simuk
+    from arviz_plots import plot_ecdf_pit
+
+    def update_data(model, augmented_data, simulation_idx):
+        with model:
+            pm.set_data({
+                "sigma_data": np.concatenate([sigma, sigma]),
+                "school_idx": np.concatenate([np.arange(8), np.arange(8)])
+            })
+
+    post_sbc = simuk.SBC(
+        centered_eight,
+        method="posterior",
+        trace=idata,
+        update_data=update_data,
+        num_simulations=50,
+        sample_kwargs={"draws": 100, "tune": 100},
+        progress_bar=False
+    )
+    post_sbc.run_simulations()
+
+    plot_ecdf_pit(post_sbc.simulations, group="posterior_sbc", visuals={"xlabel": False})
+    ```
+
+![Posterior Simulation based calibration plots, ecdf](docs/examples/img/posterior_sbc.png)
+
+We see that the funnel neck in the eight schools model is avoided and the inference algorithm is well-calibrated locally for the observed data, as indicated by the absence of red points.
 
 ## References
 

docs/conf.py

@@ -46,7 +46,7 @@
 numpydoc_show_inherited_class_members = False
 numpydoc_class_members_toctree = False
 
-source_suffix = ".rst"
+source_suffix = {".rst": "restructuredtext", ".ipynb": "myst-nb"}
 
 master_doc = "index"
 language = "en"

docs/examples.rst

@@ -7,13 +7,31 @@ The gallery below presents examples that demonstrate the use of Simuk.
    :gutter: 2 2 3 3
 
    .. grid-item-card::
-      :link: ./examples/gallery/sbc.html
+      :link: ./examples/gallery/prior_sbc.html
       :text-align: center
       :shadow: none
       :class-card: example-gallery
 
-      .. image:: examples/img/sbc.png
-         :alt: SBC
+      .. image:: examples/img/prior_sbc.png
+         :alt: Prior SBC
 
       +++
-      SBC
+      Prior SBC
+
+   .. grid-item-card::
+      :link: ./examples/gallery/posterior_sbc.html
+      :text-align: center
+      :shadow: none
+      :class-card: example-gallery
+
+      .. image:: examples/img/posterior_sbc.png
+         :alt: Posterior SBC
+      +++
+      Posterior SBC
+
+.. toctree::
+   :hidden:
+   :maxdepth: 1
+
+   examples/gallery/prior_sbc
+   examples/gallery/posterior_sbc