From e17bff6c8db447e98f408c92a7b118bf2258903c Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Mon, 10 Feb 2025 15:18:46 +0000
Subject: [PATCH 01/25] implementing sbc

---
 README.md                  |  29 --------
 forecasttools/sbc.py       | 147 +++++++++++++++++++++++++++++++++++++
 forecasttools/sbc_plots.py | 121 ++++++++++++++++++++++++++++++
 pyproject.toml             |   3 +-
 4 files changed, 270 insertions(+), 30 deletions(-)
 create mode 100644 forecasttools/sbc.py
 create mode 100644 forecasttools/sbc_plots.py

diff --git a/README.md b/README.md
index 0c0fea0..1d4c87c 100644
--- a/README.md
+++ b/README.md
@@ -5,7 +5,6 @@
 render a GitHub README, see
 <https://quarto.org/docs/output-formats/gfm.html>
 -->
-
 <!-- ```{python}
 #| echo: false
 import polars as pl
@@ -34,9 +33,7 @@ Notes:
   [pyrenew-hew](https://github.com/CDCgov/pyrenew-hew).
 
 <details>
-
 <summary>
-
 A Tentative Utilities Diagram
 </summary>
 
@@ -384,9 +381,7 @@ given the expectation that no one would ever call it; however, for
 reproducibility’s sake, the following is included here)
 
 <details>
-
 <summary>
-
 Some Of The Forecast Code
 </summary>
 
@@ -762,9 +757,7 @@ CDC use of GitHub does not imply an endorsement of any one particular
 service, product, or enterprise.
 
 <details>
-
 <summary>
-
 Rules, Policy, And Collaboration
 </summary>
 
@@ -776,14 +769,10 @@ Rules, Policy, And Collaboration
 - [Code of Conduct](./rules-and-policies/code-of-conduct.md)
 
 </details>
-
 <details>
-
 <summary>
-
 Public Domain Standard Notice
 </summary>
-
 This repository constitutes a work of the United States Government and
 is not subject to domestic copyright protection under 17 USC § 105. This
 repository is in the public domain within the United States, and
@@ -794,11 +783,8 @@ contributions to this repository will be released under the CC0
 dedication. By submitting a pull request you are agreeing to comply with
 this waiver of copyright interest.
 </details>
-
 <details>
-
 <summary>
-
 License Standard Notice
 </summary>
 
@@ -821,25 +807,18 @@ http://www.apache.org/licenses/LICENSE-2.0.html
 The source code forked from other open source projects will inherit its
 license.
 </details>
-
 <details>
-
 <summary>
-
 Privacy Standard Notice
 </summary>
-
 This repository contains only non-sensitive, publicly available data and
 information. All material and community participation is covered by the
 [Disclaimer](DISCLAIMER.md) and [Code of Conduct](code-of-conduct.md).
 For more information about CDC’s privacy policy, please visit
 [http://www.cdc.gov/other/privacy.html](https://www.cdc.gov/other/privacy.html).
 </details>
-
 <details>
-
 <summary>
-
 Contributing Standard Notice
 </summary>
 
@@ -857,27 +836,19 @@ through CDC including this GitHub page may be subject to applicable
 federal law, including but not limited to the Federal Records Act, and
 may be archived. Learn more at <http://www.cdc.gov/other/privacy.html>.
 </details>
-
 <details>
-
 <summary>
-
 Records Management Standard Notice
 </summary>
-
 This repository is not a source of government records, but is a copy to
 increase collaboration and collaborative potential. All government
 records will be published through the [CDC web
 site](http://www.cdc.gov).
 </details>
-
 <details>
-
 <summary>
-
 Additional Standard Notices
 </summary>
-
 Please refer to [CDC’s Template
 Repository](https://github.com/CDCgov/template) for more information
 about [contributing to this
diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
new file mode 100644
index 0000000..9944063
--- /dev/null
+++ b/forecasttools/sbc.py
@@ -0,0 +1,147 @@
+import logging
+
+import arviz as az
+import numpyro
+from jax import random
+from numpyro.infer import MCMC
+from numpyro.infer.mcmc import MCMCKernel
+
+
+class quiet_logging:
+    """Turn off logging for numpyro. For use in SBC."""
+
+    def __init__(self, *libraries):
+        self.loggers = [logging.getLogger(library) for library in libraries]
+
+    def __call__(self, func):
+        def wrapped(cls, *args, **kwargs):
+            levels = []
+            for logger in self.loggers:
+                levels.append(logger.level)
+                logger.setLevel(logging.CRITICAL)
+            res = func(cls, *args, **kwargs)
+            for logger, level in zip(self.loggers, levels):
+                logger.setLevel(level)
+            return res
+
+        return wrapped
+
+
+class SBC:
+    def __init__(
+        self,
+        mcmc_kernel: MCMCKernel,
+        *args,
+        observed_vars: dict[str, str],
+        num_simulations=100,
+        sample_kwargs=dict(num_warmup=500, num_samples=10),
+        seed=random.PRNGKey(1234),
+        **kwargs,
+    ):
+        """Set up class for doing SBC.
+        Based on simulation based calibration (Talts et. al. 2018) in PyMC.
+
+        Parameters
+        ----------
+        mcmc_kernel : numpyro.infer.mcmc.MCMCKernel
+            An instance of a numpyo MCMC kernel object.
+        observed_vars : dict[str, str]
+            A dictionary mapping observed/response variable name as a kwarg to the
+            numpyro model to the corresponding variable name sampled using `numpyro.sample`.
+        args: tuple
+            Positional arguments passed to `numpyro.sample`.
+        num_simulations : int
+            How many simulations to run for SBC.
+        sample_kwargs : dict[str] -> Any
+            Arguments passed to `numpyro.sample`. Defaults to `dict(num_warmup=500, num_samples=10)`.
+            Which assumes a MCMC sampler e.g. NUTS.
+        seed : random.PRNGKey
+            Random seed.
+        kwargs: dict
+            Keyword arguments passed to `numpyro` models.
+        """
+        self.mcmc_kernel = mcmc_kernel
+        if not hasattr(mcmc_kernel, 'model'):
+            raise ValueError("The `mcmc_kernel` must have a 'model' attribute.")
+        self.model = mcmc_kernel.model
+        self.args = args
+        self.kwargs = kwargs
+        self.observed_vars = observed_vars
+        self.num_simulations = num_simulations
+        self.sample_kwargs = sample_kwargs
+
+        self._simulations_complete = 0
+        prior_pred_rng, sampler_rng = random.split(seed)
+        self._prior_pred_rng = prior_pred_rng
+        self._sampler_rng = sampler_rng
+
+    def _get_prior_predictive_samples(self):
+        """Generate samples to use for the simulations."""
+        prior_predictive_fn = numpyro.infer.Predictive(self.mcmc_kernel.model, num_samples=self.num_simulations)
+        prior_predictions = prior_predictive_fn(self._prior_pred_rng, *self.args, **self.kwargs)
+        observed_values = {k: prior_predictions[v] for k, v in self.observed_vars.items()}
+        return observed_values
+
+    def _get_posterior_samples(self, seed, prior_predictive_draw):
+        """Generate posterior samples conditioned to a prior predictive sample."""
+        mcmc = MCMC(self.mcmc_kernel, **self.sample_kwargs)
+        obs_vars = {**self.kwargs, **prior_predictive_draw}
+        mcmc.run(seed, *self.args, **obs_vars)
+        return az.from_numpyro(mcmc)
+
+    # @quiet_logging("numpyro")
+    # def run_simulations(self):
+    #     """Run all the simulations.
+    #     """
+    #     observed_values = self._get_prior_predictive_samples()
+    #     sampler_seeds = random.split(self._sampler_rng, self.num_simulations)
+
+    #     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 observed_values.items()}
+
+    #             posterior = self._get_posterior_samples(sampler_seeds[idx], prior_predictive_draw)
+    #             for name in self.var_names:
+    #                 self.simulations[name].append(
+    #                     (posterior[name] < prior[name].sel(chain=0, draw=idx)).sum("sample").values
+    #                 )
+    #             self._simulations_complete += 1
+    #             progress.update()
+    #     finally:
+    #         self.simulations = {
+    #             k: v[: self._simulations_complete] for k, v in self.simulations.items()
+    #         }
+    #         progress.close()
+
+    # # def plot_results(self, kind="ecdf", var_names=None, color="C0"):
+    # #     """Visual diagnostic for SBC.
+
+    # #     Currently it support two options: `ecdf` for the empirical CDF plots
+    # #     of the difference between prior and posterior. `hist` for the rank
+    # #     histogram.
+
+
+    # #     Parameters
+    # #     ----------
+    # #     simulations : dict[str] -> listlike
+    # #         The SBC.simulations dictionary.
+    # #     kind : str
+    # #         What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
+    # #     var_names : list[str]
+    # #         Variables to plot (defaults to all)
+    # #     figsize : tuple
+    # #         Figure size for the plot. If None, it will be defined automatically.
+    # #     color : str
+    # #         Color to use for the eCDF or histogram
+
+    # #     Returns
+    # #     -------
+    # #     fig, axes
+    # #         matplotlib figure and axes
+    # #     """
+    # #     return plot_results(self.simulations, kind=kind, var_names=var_names, color=color)
diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
new file mode 100644
index 0000000..9c0e500
--- /dev/null
+++ b/forecasttools/sbc_plots.py
@@ -0,0 +1,121 @@
+"""Plots for the simulation based calibration"""
+import itertools
+
+import arviz as az
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.special import bdtrik
+
+
+def plot_results(simulations, kind="ecdf", var_names=None, figsize=None, color="C0"):
+    """Visual diagnostic for SBC.
+
+    Currently it support two options: `ecdf` for the empirical CDF plots
+    of the difference between prior and posterior. `hist` for the rank
+    histogram.
+
+
+    Parameters
+    ----------
+    simulations : dict[str] -> listlike
+        The SBC.simulations dictionary.
+    kind : str
+        What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
+    var_names : list[str]
+        Variables to plot (defaults to all)
+    figsize : tuple
+        Figure size for the plot. If None, it will be defined automatically.
+    color : str
+        Color to use for the eCDF or histogram
+
+    Returns
+    -------
+    fig, axes
+        matplotlib figure and axes
+    """
+
+    if kind not in ["ecdf", "hist"]:
+        raise ValueError(f"kind must be 'ecdf' or 'hist', not {kind}")
+
+    if var_names is None:
+        var_names = list(simulations.keys())
+
+    sims = {}
+    for k in var_names:
+        ary = np.array(simulations[k])
+        while ary.ndim < 2:
+            ary = np.expand_dims(ary, -1)
+        sims[k] = ary
+
+    n_plots = sum(np.prod(v.shape[1:]) for v in sims.values())
+
+    if n_plots > 1:
+        if figsize is None:
+            figsize=(8, n_plots*0.75)
+
+        fig, axes = plt.subplots(nrows=(n_plots + 1) // 2, ncols=2, figsize=figsize, sharex=True)
+        axes = axes.flatten()
+    else:
+        if figsize is None:
+            figsize=(8, 1.5)
+
+        fig, axes = plt.subplots(nrows=1, ncols=1, figsize=figsize)
+        axes = [axes]
+
+    if kind == "ecdf":
+        cdf = UniformCDF(len(simulations[var_names[0]]))
+
+    idx = 0
+    for var_name, var_data in sims.items():
+        plot_idxs = list(itertools.product(*(np.arange(s) for s in var_data.shape[1:])))
+        if len(plot_idxs) > 1:
+            has_dims = True
+        else:
+            has_dims = False
+
+        for indices in plot_idxs:
+            if has_dims:
+                dim_label = f'{var_name}[{"][".join(map(str, indices))}]'
+            else:
+                dim_label = var_name
+            ax = axes[idx]
+            ary = var_data[(...,) + indices]
+            if kind == "ecdf":
+                az.plot_ecdf(ary,
+                            cdf=cdf,
+                            difference = True,
+                            pit = True,
+                            confidence_bands = "auto",
+                            plot_kwargs={"color":color},
+                            fill_kwargs={"color":color},
+                            ax=ax)
+            else:
+                hist(ary, color=color, ax=ax)
+            ax.set_title(dim_label)
+            ax.set_yticks([])
+            idx += 1
+
+    for j in range(n_plots, len(axes)):
+        fig.delaxes(axes[j])
+
+    return fig, axes
+
+def hist(ary, color, ax):
+    hist, bins = np.histogram(ary, bins="auto")
+    bin_centers = 0.5 * (bins[:-1] + bins[1:])
+    max_rank = np.ceil(bins[-1])
+    len_bins = len(bins)
+    n_sims = len(ary)
+
+    band = np.ceil(bdtrik([0.025, 0.5, 0.975], n_sims, 1/len_bins))
+    ax.bar(bin_centers, hist, width=bins[1] - bins[0], color=color, edgecolor='black')
+    ax.axhline(band[1], color="0.5", ls="--")
+    ax.fill_between(np.linspace(0, max_rank, len_bins), band[0], band[2], color="0.5", alpha=0.5)
+
+
+class UniformCDF():
+    def __init__(self, upper_bound):
+        self.upper_bound = upper_bound
+
+    def __call__(self, x):
+        return np.where(x < 0, 0, np.where(x > self.upper_bound, 1, x / self.upper_bound))
diff --git a/pyproject.toml b/pyproject.toml
index 9fca911..2c214e2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,12 +31,13 @@ polars = "^1.8.2"
 xarray = "^2024.9.0"
 matplotlib = "^3.9.2"
 epiweeks = "^2.3.0"
+numpyro = "^0.17.0"
+numpy = "^2.2.2"
 
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.8.0"
 ipykernel = "^6.29.5"
-numpyro = "^0.15.3"
 diffrax = "^0.6.0"
 patsy = "^0.5.6"
 nbformat = "^5.10.4"

From 2b7e68150945f00025c1e205e8e52d4c39adc23f Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 11:23:54 +0000
Subject: [PATCH 02/25] sbc method

---
 forecasttools/sbc.py | 147 ++++++++++++++++++++++++-------------------
 pyproject.toml       |   1 +
 2 files changed, 85 insertions(+), 63 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 9944063..2974ac4 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -1,10 +1,14 @@
 import logging
 
 import arviz as az
+import jax.numpy as jnp
 import numpyro
 from jax import random
 from numpyro.infer import MCMC
 from numpyro.infer.mcmc import MCMCKernel
+from tqdm import tqdm
+
+from forecasttools.sbc_plots import plot_results
 
 
 class quiet_logging:
@@ -33,8 +37,8 @@ def __init__(
         mcmc_kernel: MCMCKernel,
         *args,
         observed_vars: dict[str, str],
-        num_simulations=100,
-        sample_kwargs=dict(num_warmup=500, num_samples=10),
+        num_simulations=10,
+        sample_kwargs=dict(num_warmup=500, num_samples=1, progress_bar = False),
         seed=random.PRNGKey(1234),
         **kwargs,
     ):
@@ -67,81 +71,98 @@ def __init__(
         self.args = args
         self.kwargs = kwargs
         self.observed_vars = observed_vars
+        for key in self.observed_vars:
+            if key in self.kwargs and self.kwargs[key] is not None:
+                raise ValueError(f"The value for '{key}' in kwargs must be None for this to be a prior predictive check.")
+
         self.num_simulations = num_simulations
         self.sample_kwargs = sample_kwargs
 
+        self.simulations = {}
         self._simulations_complete = 0
         prior_pred_rng, sampler_rng = random.split(seed)
         self._prior_pred_rng = prior_pred_rng
         self._sampler_rng = sampler_rng
 
     def _get_prior_predictive_samples(self):
-        """Generate samples to use for the simulations."""
+        """
+        Generate samples to use for the simulations by prior predictive sampling. Then splits between
+        """
         prior_predictive_fn = numpyro.infer.Predictive(self.mcmc_kernel.model, num_samples=self.num_simulations)
         prior_predictions = prior_predictive_fn(self._prior_pred_rng, *self.args, **self.kwargs)
-        observed_values = {k: prior_predictions[v] for k, v in self.observed_vars.items()}
-        return observed_values
+        prior_pred = {k: prior_predictions[v] for k, v in self.observed_vars.items()}
+        prior = {k: v for k, v in prior_predictions.items() if k not in self.observed_vars.values()}
+        return prior, prior_pred
 
     def _get_posterior_samples(self, seed, prior_predictive_draw):
         """Generate posterior samples conditioned to a prior predictive sample."""
         mcmc = MCMC(self.mcmc_kernel, **self.sample_kwargs)
         obs_vars = {**self.kwargs, **prior_predictive_draw}
         mcmc.run(seed, *self.args, **obs_vars)
-        return az.from_numpyro(mcmc)
-
-    # @quiet_logging("numpyro")
-    # def run_simulations(self):
-    #     """Run all the simulations.
-    #     """
-    #     observed_values = self._get_prior_predictive_samples()
-    #     sampler_seeds = random.split(self._sampler_rng, self.num_simulations)
-
-    #     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 observed_values.items()}
-
-    #             posterior = self._get_posterior_samples(sampler_seeds[idx], prior_predictive_draw)
-    #             for name in self.var_names:
-    #                 self.simulations[name].append(
-    #                     (posterior[name] < prior[name].sel(chain=0, draw=idx)).sum("sample").values
-    #                 )
-    #             self._simulations_complete += 1
-    #             progress.update()
-    #     finally:
-    #         self.simulations = {
-    #             k: v[: self._simulations_complete] for k, v in self.simulations.items()
-    #         }
-    #         progress.close()
-
-    # # def plot_results(self, kind="ecdf", var_names=None, color="C0"):
-    # #     """Visual diagnostic for SBC.
-
-    # #     Currently it support two options: `ecdf` for the empirical CDF plots
-    # #     of the difference between prior and posterior. `hist` for the rank
-    # #     histogram.
-
-
-    # #     Parameters
-    # #     ----------
-    # #     simulations : dict[str] -> listlike
-    # #         The SBC.simulations dictionary.
-    # #     kind : str
-    # #         What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
-    # #     var_names : list[str]
-    # #         Variables to plot (defaults to all)
-    # #     figsize : tuple
-    # #         Figure size for the plot. If None, it will be defined automatically.
-    # #     color : str
-    # #         Color to use for the eCDF or histogram
-
-    # #     Returns
-    # #     -------
-    # #     fig, axes
-    # #         matplotlib figure and axes
-    # #     """
-    # #     return plot_results(self.simulations, kind=kind, var_names=var_names, color=color)
+        idata = az.from_numpyro(mcmc)
+        return idata
+
+    @quiet_logging("numpyro")
+    def run_simulations(self):
+        """
+        Run all the simulations.
+        """
+        prior, prior_pred = self._get_prior_predictive_samples()
+        sampler_seeds = random.split(self._sampler_rng, self.num_simulations)
+        self.simulations = {name: [] for name in prior.keys()}
+        progress = tqdm(
+            initial=self._simulations_complete,
+            total=self.num_simulations,
+        )
+        try:
+            while self._simulations_complete < self.num_simulations:
+                idx = self._simulations_complete
+                prior_draw = {k:v[idx] for k, v in prior.items()}
+                prior_predictive_draw = {k:v[idx] for k, v in prior_pred.items()}
+                posterior = self._get_posterior_samples(sampler_seeds[idx], prior_predictive_draw)['posterior']
+                # pdb.set_trace()
+                for name in prior.keys():
+                    num_dims = jnp.ndim(prior_draw[name])
+                    if num_dims == 0:
+                        self.simulations[name].append(
+                            (posterior[name].sel(chain=0) < prior_draw[name]).sum().values
+                        )
+                    else:
+                        self.simulations[name].append(
+                            (posterior[name].sel(chain=0) < prior_draw[name]).sum(axis=0).values
+                        )
+                self._simulations_complete += 1
+                progress.update()
+        finally:
+            self.simulations = {
+                k: v[: self._simulations_complete] for k, v in self.simulations.items()
+            }
+            progress.close()
+
+    def plot_results(self, kind="ecdf", var_names=None, color="C0"):
+        """Visual diagnostic for SBC.
+
+        Currently it support two options: `ecdf` for the empirical CDF plots
+        of the difference between prior and posterior. `hist` for the rank
+        histogram.
+
+
+        Parameters
+        ----------
+        simulations : dict[str] -> listlike
+            The SBC.simulations dictionary.
+        kind : str
+            What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
+        var_names : list[str]
+            Variables to plot (defaults to all)
+        figsize : tuple
+            Figure size for the plot. If None, it will be defined automatically.
+        color : str
+            Color to use for the eCDF or histogram
+
+        Returns
+        -------
+        fig, axes
+            matplotlib figure and axes
+        """
+        return plot_results(self.simulations, kind=kind, var_names=var_names, color=color)
diff --git a/pyproject.toml b/pyproject.toml
index 2c214e2..cbb8b7e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,6 +33,7 @@ matplotlib = "^3.9.2"
 epiweeks = "^2.3.0"
 numpyro = "^0.17.0"
 numpy = "^2.2.2"
+tqdm = "^4.67.1"
 
 
 [tool.poetry.group.dev.dependencies]

From fda44bc62dce69fb83b1423a39b7eebb3bc9f7a5 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 11:41:03 +0000
Subject: [PATCH 03/25] pass number of draws to plotting

---
 forecasttools/sbc.py       | 37 +++++++++++--------------------------
 forecasttools/sbc_plots.py |  8 +++++---
 2 files changed, 16 insertions(+), 29 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 2974ac4..8edfd53 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -1,4 +1,3 @@
-import logging
 
 import arviz as az
 import jax.numpy as jnp
@@ -11,26 +10,6 @@
 from forecasttools.sbc_plots import plot_results
 
 
-class quiet_logging:
-    """Turn off logging for numpyro. For use in SBC."""
-
-    def __init__(self, *libraries):
-        self.loggers = [logging.getLogger(library) for library in libraries]
-
-    def __call__(self, func):
-        def wrapped(cls, *args, **kwargs):
-            levels = []
-            for logger in self.loggers:
-                levels.append(logger.level)
-                logger.setLevel(logging.CRITICAL)
-            res = func(cls, *args, **kwargs)
-            for logger, level in zip(self.loggers, levels):
-                logger.setLevel(level)
-            return res
-
-        return wrapped
-
-
 class SBC:
     def __init__(
         self,
@@ -38,7 +17,7 @@ def __init__(
         *args,
         observed_vars: dict[str, str],
         num_simulations=10,
-        sample_kwargs=dict(num_warmup=500, num_samples=1, progress_bar = False),
+        sample_kwargs=dict(num_warmup=500, num_samples=100, progress_bar = False),
         seed=random.PRNGKey(1234),
         **kwargs,
     ):
@@ -83,6 +62,7 @@ def __init__(
         prior_pred_rng, sampler_rng = random.split(seed)
         self._prior_pred_rng = prior_pred_rng
         self._sampler_rng = sampler_rng
+        self.num_samples = None
 
     def _get_prior_predictive_samples(self):
         """
@@ -99,10 +79,10 @@ def _get_posterior_samples(self, seed, prior_predictive_draw):
         mcmc = MCMC(self.mcmc_kernel, **self.sample_kwargs)
         obs_vars = {**self.kwargs, **prior_predictive_draw}
         mcmc.run(seed, *self.args, **obs_vars)
+        num_samples = mcmc.num_samples
         idata = az.from_numpyro(mcmc)
-        return idata
+        return idata, num_samples
 
-    @quiet_logging("numpyro")
     def run_simulations(self):
         """
         Run all the simulations.
@@ -119,7 +99,12 @@ def run_simulations(self):
                 idx = self._simulations_complete
                 prior_draw = {k:v[idx] for k, v in prior.items()}
                 prior_predictive_draw = {k:v[idx] for k, v in prior_pred.items()}
-                posterior = self._get_posterior_samples(sampler_seeds[idx], prior_predictive_draw)['posterior']
+                idata, num_samples = self._get_posterior_samples(sampler_seeds[idx], prior_predictive_draw)
+                if self.num_samples is None:
+                    self.num_samples = num_samples
+                if self.num_samples != num_samples:
+                    raise ValueError("The number of samples from the posterior is not consistent.")
+                posterior = idata['posterior']
                 # pdb.set_trace()
                 for name in prior.keys():
                     num_dims = jnp.ndim(prior_draw[name])
@@ -165,4 +150,4 @@ def plot_results(self, kind="ecdf", var_names=None, color="C0"):
         fig, axes
             matplotlib figure and axes
         """
-        return plot_results(self.simulations, kind=kind, var_names=var_names, color=color)
+        return plot_results(self.simulations, self.num_samples, kind=kind, var_names=var_names, color=color)
diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index 9c0e500..f7cb3ec 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -7,7 +7,7 @@
 from scipy.special import bdtrik
 
 
-def plot_results(simulations, kind="ecdf", var_names=None, figsize=None, color="C0"):
+def plot_results(simulations, ndraws, kind="ecdf", var_names=None, figsize=None, color="C0"):
     """Visual diagnostic for SBC.
 
     Currently it support two options: `ecdf` for the empirical CDF plots
@@ -19,6 +19,8 @@ def plot_results(simulations, kind="ecdf", var_names=None, figsize=None, color="
     ----------
     simulations : dict[str] -> listlike
         The SBC.simulations dictionary.
+    ndraws : int
+        Number of draws in each posterior predictive sample
     kind : str
         What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
     var_names : list[str]
@@ -51,7 +53,7 @@ def plot_results(simulations, kind="ecdf", var_names=None, figsize=None, color="
 
     if n_plots > 1:
         if figsize is None:
-            figsize=(8, n_plots*0.75)
+            figsize=(8, n_plots*1.0)
 
         fig, axes = plt.subplots(nrows=(n_plots + 1) // 2, ncols=2, figsize=figsize, sharex=True)
         axes = axes.flatten()
@@ -63,7 +65,7 @@ def plot_results(simulations, kind="ecdf", var_names=None, figsize=None, color="
         axes = [axes]
 
     if kind == "ecdf":
-        cdf = UniformCDF(len(simulations[var_names[0]]))
+        cdf = UniformCDF(ndraws)
 
     idx = 0
     for var_name, var_data in sims.items():

From 66fb72e29c522775ef90d0db02fda0aa50e0da01 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 11:47:52 +0000
Subject: [PATCH 04/25] Update sbc.py

---
 forecasttools/sbc.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 8edfd53..5773be9 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -31,25 +31,27 @@ def __init__(
         observed_vars : dict[str, str]
             A dictionary mapping observed/response variable name as a kwarg to the
             numpyro model to the corresponding variable name sampled using `numpyro.sample`.
-        args: tuple
+        args : tuple
             Positional arguments passed to `numpyro.sample`.
         num_simulations : int
             How many simulations to run for SBC.
         sample_kwargs : dict[str] -> Any
-            Arguments passed to `numpyro.sample`. Defaults to `dict(num_warmup=500, num_samples=10)`.
+            Arguments passed to `numpyro.sample`. Defaults to `dict(num_warmup=500, num_samples=100, progress_bar = False)`.
             Which assumes a MCMC sampler e.g. NUTS.
         seed : random.PRNGKey
             Random seed.
-        kwargs: dict
+        kwargs : dict
             Keyword arguments passed to `numpyro` models.
         """
         self.mcmc_kernel = mcmc_kernel
         if not hasattr(mcmc_kernel, 'model'):
             raise ValueError("The `mcmc_kernel` must have a 'model' attribute.")
+
         self.model = mcmc_kernel.model
         self.args = args
         self.kwargs = kwargs
         self.observed_vars = observed_vars
+
         for key in self.observed_vars:
             if key in self.kwargs and self.kwargs[key] is not None:
                 raise ValueError(f"The value for '{key}' in kwargs must be None for this to be a prior predictive check.")
@@ -67,6 +69,7 @@ def __init__(
     def _get_prior_predictive_samples(self):
         """
         Generate samples to use for the simulations by prior predictive sampling. Then splits between
+        observed and unobserved variables based on the `observed_vars` attribute.
         """
         prior_predictive_fn = numpyro.infer.Predictive(self.mcmc_kernel.model, num_samples=self.num_simulations)
         prior_predictions = prior_predictive_fn(self._prior_pred_rng, *self.args, **self.kwargs)
@@ -75,7 +78,11 @@ def _get_prior_predictive_samples(self):
         return prior, prior_pred
 
     def _get_posterior_samples(self, seed, prior_predictive_draw):
-        """Generate posterior samples conditioned to a prior predictive sample."""
+        """
+        Generate posterior samples conditioned to a prior predictive sample.
+        This returns the posterior samples and the number of samples. The number of samples are used
+        in scaling plotting and checking that each inference draw has the same number of samples.
+        """
         mcmc = MCMC(self.mcmc_kernel, **self.sample_kwargs)
         obs_vars = {**self.kwargs, **prior_predictive_draw}
         mcmc.run(seed, *self.args, **obs_vars)
@@ -85,7 +92,8 @@ def _get_posterior_samples(self, seed, prior_predictive_draw):
 
     def run_simulations(self):
         """
-        Run all the simulations.
+        The main method of `SBC` class that runs the simulations for simulation based calibration and
+        fills the `simulations` attribute with the results.
         """
         prior, prior_pred = self._get_prior_predictive_samples()
         sampler_seeds = random.split(self._sampler_rng, self.num_simulations)
@@ -134,7 +142,7 @@ def plot_results(self, kind="ecdf", var_names=None, color="C0"):
 
         Parameters
         ----------
-        simulations : dict[str] -> listlike
+        simulations
             The SBC.simulations dictionary.
         kind : str
             What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'

From fa45689b752dddaf7989d55e73151651b50a8cc1 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 12:13:05 +0000
Subject: [PATCH 05/25] lint

---
 forecasttools/sbc.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 5773be9..457a3b2 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -1,4 +1,3 @@
-
 import arviz as az
 import jax.numpy as jnp
 import numpyro

From e8b46483b3412837997b05f36de4ad5f421c7481 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 12:16:29 +0000
Subject: [PATCH 06/25] unit tests

---
 tests/test_sbc.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 tests/test_sbc.py

diff --git a/tests/test_sbc.py b/tests/test_sbc.py
new file mode 100644
index 0000000..5bdc2dd
--- /dev/null
+++ b/tests/test_sbc.py
@@ -0,0 +1,71 @@
+import numpyro
+import pytest
+from jax import random
+from numpyro.infer import NUTS
+
+from forecasttools.sbc import SBC
+
+
+@pytest.fixture
+def simple_model():
+    def model(y=None):
+        mu = numpyro.sample('mu', numpyro.distributions.Normal(0, 1))
+        numpyro.sample('z', numpyro.distributions.Normal(mu, 1), obs=y)
+    return model
+
+@pytest.fixture
+def mcmc_kernel(simple_model):
+    return NUTS(simple_model)
+
+@pytest.fixture
+def observed_vars():
+    return {"y" : "z"}
+
+@pytest.fixture
+def sbc_instance(mcmc_kernel, observed_vars):
+    return SBC(mcmc_kernel, y = None, observed_vars=observed_vars)
+
+def test_sbc_initialization(sbc_instance, mcmc_kernel, observed_vars):
+    """
+    Test that the SBC class is initialized correctly.
+    """
+    assert sbc_instance.mcmc_kernel == mcmc_kernel
+    assert sbc_instance.observed_vars == observed_vars
+    assert sbc_instance.num_simulations == 10
+    assert sbc_instance.sample_kwargs == dict(num_warmup=500, num_samples=100, progress_bar=False)
+    assert sbc_instance._simulations_complete == 0
+
+def test_get_prior_predictive_samples(sbc_instance):
+    """
+    Test that the prior and prior predictive samples are generated correctly.
+    """
+    prior, prior_pred = sbc_instance._get_prior_predictive_samples()
+    assert "y" in prior_pred
+    assert "mu" in prior
+
+def test_get_posterior_samples(sbc_instance):
+    """
+    Test that the posterior samples are generated correctly.
+    """
+    prior, prior_pred = sbc_instance._get_prior_predictive_samples()
+    seed = random.PRNGKey(0)
+    idata, num_samples = sbc_instance._get_posterior_samples(seed, prior_pred)
+    assert "posterior" in idata
+    assert num_samples == 100
+
+def test_run_simulations(sbc_instance):
+    """
+    Test that the simulations for SBC are run correctly.
+    """
+    sbc_instance.run_simulations()
+    assert sbc_instance._simulations_complete == sbc_instance.num_simulations
+    assert "mu" in sbc_instance.simulations
+
+def test_plot_results(sbc_instance):
+    """
+    Test that the results are plotted.
+    """
+    sbc_instance.run_simulations()
+    fig, axes = sbc_instance.plot_results()
+    assert fig is not None
+    assert axes is not None

From 0f38bc20113693ff68eaa4437f5efa0609192c3a Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 14:25:59 +0000
Subject: [PATCH 07/25] Tidying up and extending example

---
 forecasttools/sbc.py             |   1 -
 notebooks/sbc_model_checking.qmd | 139 +++++++++++++++++++++++++++++++
 pyproject.toml                   |   1 +
 tests/test_sbc.py                |  11 +++
 4 files changed, 151 insertions(+), 1 deletion(-)
 create mode 100644 notebooks/sbc_model_checking.qmd

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 457a3b2..0328550 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -112,7 +112,6 @@ def run_simulations(self):
                 if self.num_samples != num_samples:
                     raise ValueError("The number of samples from the posterior is not consistent.")
                 posterior = idata['posterior']
-                # pdb.set_trace()
                 for name in prior.keys():
                     num_dims = jnp.ndim(prior_draw[name])
                     if num_dims == 0:
diff --git a/notebooks/sbc_model_checking.qmd b/notebooks/sbc_model_checking.qmd
new file mode 100644
index 0000000..91328aa
--- /dev/null
+++ b/notebooks/sbc_model_checking.qmd
@@ -0,0 +1,139 @@
+---
+title: "Using the `SBC` class to do sim-based calibration for `numpyro` models"
+format: gfm
+engine: jupyter
+execute:
+    warning: false
+    output: false
+    cache: true
+bibliography: references.bib
+---
+
+This notebook covers:
+
+-   A brief introduction to simulation-based calibration (SBC) as a method for testing the self-consistency of a Bayesian inference method.
+-   Tuning and running SBC for a `numpyro` model using the `SBC` from `forecasttools`.
+-   Plotting results and interpreting them.
+
+The example we will use to illustrate our ideas is the classic "eight schools" inference problem from section 5.5 of Gelman *et al* [@gelman2013] with posterior sampling done using the NUTS algorithm implemented by `numpyro` PPL.
+
+## Simulation-based calibration
+
+"Unit" testing the correctness of Bayesian probabilistic programs is challenging because the target quantity is a distribution (the posterior distribution) typically sampled from via a random process, e.g. Markov chain Monte Carlo sampling. The range of models where the correct posterior distribution is known analytically, and therefore, we can test sampling against distributional checks (e.g. Kolmogorov-Smirnov tests, chi-squared tests, etc) is limited.
+
+Simulation-based calibration (SBC) [@talts] aims to provide a correctness check for a Bayesian model and sampler using the self-consistency property of Bayesian inference:
+
+$$
+p(\theta) = \int p(y,\theta) dy = \int p(\theta | y) p(y) dy = \int p(\theta|y)\int p(y | \theta')p(\theta')d\theta' dy.
+$$
+
+In words, this states that the distribution of posterior samples from Bayesian inference gathered over many datasets $(y^{(i)})_{i=1, 2, \dots}$ is the prior distribution $p(\theta)$ *if* the datasets are generated using the prior distribution $y^{(i)} \sim p(y | \theta') p(\theta')$.
+
+This self-consistency property of Bayesian inference provides a convenient method for checking if the model inference process $p(\theta | y)$ is performing as desired, since generating datasets $(y^{(i)})_{i=1, 2, \dots}$ from a generative model is typically quite easy. Since for any individual $k$th parameter $\theta[k] \in \theta$ we know the prior distribution $p(\theta[k])$ we can check correctness against:
+
+$$
+F^{(-1)}(\theta[k]) \sim \mathcal{U}(0,1).
+$$
+
+Where $F^{(-1)}$ is the (pseudo) inverse distribution function of the known distribution $p(\theta[k])$. This is often call ed the probability integral transform (PIT). Significant (in a classical statistics sense) deviations in the PIT of parameters from the uniform distribution form evidence that either the model is incorrectly coded, its priors are cover hard to sample regimes, and/or the Bayesian inference method is approximate/invalid. See Talts *et al* [@talts] for more details on understanding the returned tests.
+
+More precisely, in SBC:
+
+1.  For $i = 1,\dots,n$: sample the observable and parameters from the model using the prior distribution, $(y^{(i)}, \theta^{(i)})$.
+2.  For $i = 1,\dots,n$: sample posterior parameters using the Bayesian method and model under evaluation, $\theta_p^{(i)} \sim p(\theta_p | y^{(i)})$.
+3.  Generate the probability integral transform (PIT) of each of $k = 1,\dots,P$ parameters $\theta[k]$ with respect to the known prior distribution.
+4.  Assess the PIT distributions for deviation against an assumption of being $\mathcal{U}([0,1])$.
+
+In `SBC` we from the PIT by looking at the distribution of proportion of posterior samples that are less than the "true" parameter values cached along with the generated data $y^{(i)}$: $P(\theta_p^{(i)}[k] < \theta^{(i)}[k])$. This is a more convenient approach in `numpyro` than trying to solve the inverse distribution function directly.
+
+## Example: Eight schools
+
+The eight schools example is a classic example of using partial pooling to share inferential strength between groups, cf Gelman *et al* [@gelman2013]:
+
+> *A study was performed for the Educational Testing Service to analyze the effects of special coaching programs for SAT-V (Scholastic Aptitude Test-Verbal) in each of eight high schools. The outcome variable in each study was the score on a special administration of the SAT-V, a standardized multiple choice test administered by the Educational Testing Service and used to help colleges make admissions decisions; the scores can vary between 200 and 800, with mean about 500 and standard deviation about 100. The SAT examinations are designed to be resistant to short-term efforts directed specifically toward improving performance on the test; instead they are designed to reflect knowledge acquired and abilities developed over many years of education. Nevertheless, each of the eight schools in this study considered its short-term coaching program to be very successful at increasing SAT scores. Also, there was no prior reason to believe that any of the eight programs was more effective than any other or that some were more similar in effect to each other than to any other.*
+
+The statistical model for the SAT scores in each of the $J=8$ schools $y_j$ is:
+
+$$
+\begin{aligned}
+\mu & \sim \mathcal{N}(0, 5), \\
+\tau & \sim \text{HalfCauchy}(5),\\
+\theta_j & \sim \mathcal{N}(\mu, \tau),~ j = 1,\dots,J, \\
+y_j & \sim \mathcal{N}(\theta_j,\sigma_j),~ j = 1,\dots,J.
+\end{aligned}
+$$
+
+Where the the SAT standard deviations per high school $\sigma_j$ are treated as known along with the scores.
+
+Gelman *et al* use the eight schools example to illustrate partial pooling, and to demonstrate the importance of choosing the variance priors carefully.
+
+We start by setting the dependencies and the basic data and random seed.
+
+```{python}
+# Dependencies
+import arviz as az
+import numpyro
+import jax.numpy as jnp
+import numpyro.distributions as dist
+from jax import random
+from numpyro.infer import NUTS
+
+import forecasttools.sbc as sbc
+
+```
+
+```{python}
+rng_key = random.PRNGKey(0)
+J = 8
+y = jnp.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
+sigma = jnp.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
+```
+
+### Running SBC for example with HalfCauchy priors
+
+Next we define the model above as a `numpyro` model.
+
+```{python}
+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('obs', dist.Normal(theta, sigma), obs=y)
+
+nuts_kernel_cauchy_prior = NUTS(eight_schools_cauchy_prior)
+```
+
+Note that we define a `numpyro` `MCMCKernel` object to pass to `SBC`. This wraps the *model* with the *sampling approach*, which are both checked by SBC.
+
+Next, we create an instance of the `SBC` class. There are a few things to note here:
+
+-   The structure here is that we must pass the `MCMCKernel` object as well as the positional and keyword arguments that define a specific model, for example the `J` and `sigma` arguments. This follows a standard pattern in `numpyro`.
+-   Note that we don't pass the data argument `y` (and indeed this will throw an error). The reason is that we will be generating datasets for SBC rather than using any particular observed dataset.
+-   `observed_vars` is an important argument which lets `SBC` know what variables to treat as observed in the SBC process. It is a dictionary rather than a list because it maps the data name *as it is passed to the* `numpyro` *model* to the random variable name inside the `numpyro` model (which can be distinct). For example, in `eight_schools_cauchy_prior` the data value of the observed SAT scores is called `y` but the random sample is traced as `obs` . Hence, we must pass `observed_vars = dict(y = "obs")` .
+-   `sample_kwargs` gets passed through to the sampler. Note that `progress_bar = False` is important if we don't want to see alot of progress bars.
+-   `num_simulations` sets the number of SBC trials to run.
+
+```{python}
+seed1, seed2 = random.split(rng_key)
+S = sbc.SBC(nuts_kernel_cauchy_prior, J, sigma,
+    observed_vars = dict(y = "obs"),
+    sample_kwargs=dict(num_warmup=500, num_samples=1000, progress_bar = False),
+    num_simulations=100,
+    seed = seed1)
+```
+
+`SBC` class instances have methods for running the SBC simulations and plotting the results as a histogram (against binomial sampling within a bin under the uniform target distribution) and as an empirical CDF function in PIT model (this is from the `arviz` plotting utilities).
+
+```{python}
+S.run_simulations()
+```
+
+```{python}
+#|
+S.plot_results(kind = "hist")
+```
+
+```{python}
+S.plot_results()
+```
diff --git a/pyproject.toml b/pyproject.toml
index cbb8b7e..df11d1f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,6 +44,7 @@ patsy = "^0.5.6"
 nbformat = "^5.10.4"
 nbclient = "^0.10.0"
 jupyter = "^1.1.1"
+jupyter-cache = "^1.0.1"
 
 
 [tool.poetry.group.test.dependencies]
diff --git a/tests/test_sbc.py b/tests/test_sbc.py
index 5bdc2dd..b7e690c 100644
--- a/tests/test_sbc.py
+++ b/tests/test_sbc.py
@@ -1,3 +1,14 @@
+"""
+Test the SBC class using a simple model.
+
+```math
+\begin{aligned}
+\mu &\sim \text{Normal}(0, 1), \\
+z &\sim \text{Normal}(\mu, 1).
+\end{aligned}
+```
+"""
+
 import numpyro
 import pytest
 from jax import random

From cb6e2b55d23c38ab0ec8c7ece8f90b9d8d3e75b7 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Tue, 11 Feb 2025 14:43:21 +0000
Subject: [PATCH 08/25] update note

---
 notebooks/references.bib         | 16 +++++++++++
 notebooks/sbc_model_checking.qmd | 47 +++++++++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 notebooks/references.bib

diff --git a/notebooks/references.bib b/notebooks/references.bib
new file mode 100644
index 0000000..f347b61
--- /dev/null
+++ b/notebooks/references.bib
@@ -0,0 +1,16 @@
+
+@book{gelman2013,
+	title = {Bayesian Data Analysis},
+	author = {Gelman, Andrew and Carlin, John B and Stern, Hal S and Dunson, David B and Vehtari, Aki and Rubin, Donald B},
+	year = {2013},
+	date = {2013},
+	publisher = {CRC Press},
+	edition = {third},
+	langid = {en}
+}
+
+@article{talts,
+	title = {Validating Bayesian Inference Algorithms with Simulation-Based Calibration},
+	author = {Talts, Sean and Betancourt, Michael and Simpson, Daniel and Vehtari, Aki and Gelman, Andrew},
+	langid = {en}
+}
diff --git a/notebooks/sbc_model_checking.qmd b/notebooks/sbc_model_checking.qmd
index 91328aa..4522595 100644
--- a/notebooks/sbc_model_checking.qmd
+++ b/notebooks/sbc_model_checking.qmd
@@ -130,10 +130,55 @@ S.run_simulations()
 ```
 
 ```{python}
-#|
+#| output: true
+#| fig-cap: "Histogram plot of SBC results"
 S.plot_results(kind = "hist")
 ```
 
 ```{python}
+#| output: true
+#| fig-cap: "ECDF/PIT plot of SBC results"
 S.plot_results()
 ```
+
+We can see that the SBC has identified a problem with the model. The Bayesian inference looks to be slightly, but systematically, over-estimating the inter-school variation parameter $\tau$ compared to simulated values.
+
+Lets use a more informative prior from [this implementation](https://www.tensorflow.org/probability/examples/Eight_Schools).
+
+$$
+\tau \sim \text{LogNormal}(5, 1).
+$$
+
+```{python}
+def eight_schools_lognormal_prior(J, sigma, y=None):
+    mu = numpyro.sample('mu', dist.Normal(0, 5))
+    tau = numpyro.sample('tau', dist.LogNormal(5, 1))
+    with numpyro.plate('J', J):
+        theta = numpyro.sample('theta', dist.Normal(mu, tau))
+    numpyro.sample('obs', dist.Normal(theta, sigma), obs=y)
+
+nuts_kernel_lognormal_prior = NUTS(eight_schools_lognormal_prior)
+
+S2 = sbc.SBC(nuts_kernel_lognormal_prior, J, sigma,
+    observed_vars = dict(y = "obs"),
+    sample_kwargs=dict(num_warmup=500, num_samples=1000, progress_bar = False),
+    num_simulations=100,
+    seed = seed2)
+S2.run_simulations()
+```
+
+```{python}
+#| output: true
+#| fig-cap: "Histogram plot of SBC results"
+S2.plot_results(kind = "hist")
+```
+
+```{python}
+#| output: true
+#| fig-cap: "ECDF/PIT plot of SBC results"
+S2.plot_results()
+```
+
+The results of this SBC approach still indicate potential problems with self-consistent identification of every high school specific treatment effect but the self-consistent ability to infer the population parameters $\mu$ and $\tau$ seems to now work.
+
+## References

From fb7b16d57434231df6e2f497d10d6b54edd8842e Mon Sep 17 00:00:00 2001
From: AFg6K7h4fhy2 <127630341+AFg6K7h4fhy2@users.noreply.github.com>
Date: Tue, 11 Feb 2025 13:43:28 -0500
Subject: [PATCH 09/25] pre-commit fixes

---
 forecasttools/sbc.py       | 110 ++++++++++++++++++++++++++-----------
 forecasttools/sbc_plots.py |  71 +++++++++++++++---------
 tests/test_sbc.py          |  27 ++++++---
 3 files changed, 143 insertions(+), 65 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 0328550..da02fa3 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -16,8 +16,8 @@ def __init__(
         *args,
         observed_vars: dict[str, str],
         num_simulations=10,
-        sample_kwargs=dict(num_warmup=500, num_samples=100, progress_bar = False),
-        seed=random.PRNGKey(1234),
+        sample_kwargs=None,
+        seed=1234,
         **kwargs,
     ):
         """Set up class for doing SBC.
@@ -28,23 +28,33 @@ def __init__(
         mcmc_kernel : numpyro.infer.mcmc.MCMCKernel
             An instance of a numpyo MCMC kernel object.
         observed_vars : dict[str, str]
-            A dictionary mapping observed/response variable name as a kwarg to the
-            numpyro model to the corresponding variable name sampled using `numpyro.sample`.
+            A dictionary mapping observed/response variable name as a kwarg to
+            the numpyro model to the corresponding variable name sampled using
+            `numpyro.sample`.
         args : tuple
             Positional arguments passed to `numpyro.sample`.
         num_simulations : int
             How many simulations to run for SBC.
         sample_kwargs : dict[str] -> Any
-            Arguments passed to `numpyro.sample`. Defaults to `dict(num_warmup=500, num_samples=100, progress_bar = False)`.
+            Arguments passed to `numpyro.sample`. Defaults to
+            `dict(num_warmup=500, num_samples=100, progress_bar = False)`.
             Which assumes a MCMC sampler e.g. NUTS.
-        seed : random.PRNGKey
+        seed : int
             Random seed.
         kwargs : dict
             Keyword arguments passed to `numpyro` models.
         """
+        if sample_kwargs is None:
+            sample_kwargs = dict(
+                num_warmup=500, num_samples=100, progress_bar=False
+            )
+        seed = random.PRNGKey(seed)
+
         self.mcmc_kernel = mcmc_kernel
-        if not hasattr(mcmc_kernel, 'model'):
-            raise ValueError("The `mcmc_kernel` must have a 'model' attribute.")
+        if not hasattr(mcmc_kernel, "model"):
+            raise ValueError(
+                "The `mcmc_kernel` must have a 'model' attribute."
+            )
 
         self.model = mcmc_kernel.model
         self.args = args
@@ -53,7 +63,10 @@ def __init__(
 
         for key in self.observed_vars:
             if key in self.kwargs and self.kwargs[key] is not None:
-                raise ValueError(f"The value for '{key}' in kwargs must be None for this to be a prior predictive check.")
+                raise ValueError(
+                    f"The value for '{key}' in kwargs must be None for this to"
+                    " be a prior predictive check."
+                )
 
         self.num_simulations = num_simulations
         self.sample_kwargs = sample_kwargs
@@ -67,20 +80,32 @@ def __init__(
 
     def _get_prior_predictive_samples(self):
         """
-        Generate samples to use for the simulations by prior predictive sampling. Then splits between
-        observed and unobserved variables based on the `observed_vars` attribute.
+        Generate samples to use for the simulations by prior predictive
+        sampling. Then splits between observed and unobserved variables based
+        on the `observed_vars` attribute.
         """
-        prior_predictive_fn = numpyro.infer.Predictive(self.mcmc_kernel.model, num_samples=self.num_simulations)
-        prior_predictions = prior_predictive_fn(self._prior_pred_rng, *self.args, **self.kwargs)
-        prior_pred = {k: prior_predictions[v] for k, v in self.observed_vars.items()}
-        prior = {k: v for k, v in prior_predictions.items() if k not in self.observed_vars.values()}
+        prior_predictive_fn = numpyro.infer.Predictive(
+            self.mcmc_kernel.model, num_samples=self.num_simulations
+        )
+        prior_predictions = prior_predictive_fn(
+            self._prior_pred_rng, *self.args, **self.kwargs
+        )
+        prior_pred = {
+            k: prior_predictions[v] for k, v in self.observed_vars.items()
+        }
+        prior = {
+            k: v
+            for k, v in prior_predictions.items()
+            if k not in self.observed_vars.values()
+        }
         return prior, prior_pred
 
     def _get_posterior_samples(self, seed, prior_predictive_draw):
         """
         Generate posterior samples conditioned to a prior predictive sample.
-        This returns the posterior samples and the number of samples. The number of samples are used
-        in scaling plotting and checking that each inference draw has the same number of samples.
+        This returns the posterior samples and the number of samples. The
+        number of samples are used in scaling plotting and checking that each
+        inference draw has the same number of samples.
         """
         mcmc = MCMC(self.mcmc_kernel, **self.sample_kwargs)
         obs_vars = {**self.kwargs, **prior_predictive_draw}
@@ -91,12 +116,13 @@ def _get_posterior_samples(self, seed, prior_predictive_draw):
 
     def run_simulations(self):
         """
-        The main method of `SBC` class that runs the simulations for simulation based calibration and
-        fills the `simulations` attribute with the results.
+        The main method of `SBC` class that runs the simulations for
+        simulation based calibration and fills the `simulations` attribute
+        with the results.
         """
         prior, prior_pred = self._get_prior_predictive_samples()
         sampler_seeds = random.split(self._sampler_rng, self.num_simulations)
-        self.simulations = {name: [] for name in prior.keys()}
+        self.simulations = {name: [] for name in prior}
         progress = tqdm(
             initial=self._simulations_complete,
             total=self.num_simulations,
@@ -104,29 +130,41 @@ def run_simulations(self):
         try:
             while self._simulations_complete < self.num_simulations:
                 idx = self._simulations_complete
-                prior_draw = {k:v[idx] for k, v in prior.items()}
-                prior_predictive_draw = {k:v[idx] for k, v in prior_pred.items()}
-                idata, num_samples = self._get_posterior_samples(sampler_seeds[idx], prior_predictive_draw)
+                prior_draw = {k: v[idx] for k, v in prior.items()}
+                prior_predictive_draw = {
+                    k: v[idx] for k, v in prior_pred.items()
+                }
+                idata, num_samples = self._get_posterior_samples(
+                    sampler_seeds[idx], prior_predictive_draw
+                )
                 if self.num_samples is None:
                     self.num_samples = num_samples
                 if self.num_samples != num_samples:
-                    raise ValueError("The number of samples from the posterior is not consistent.")
-                posterior = idata['posterior']
-                for name in prior.keys():
+                    raise ValueError(
+                        "The number of samples from the posterior is not"
+                        " consistent."
+                    )
+                posterior = idata["posterior"]
+                for name in prior:
                     num_dims = jnp.ndim(prior_draw[name])
                     if num_dims == 0:
                         self.simulations[name].append(
-                            (posterior[name].sel(chain=0) < prior_draw[name]).sum().values
+                            (posterior[name].sel(chain=0) < prior_draw[name])
+                            .sum()
+                            .values
                         )
                     else:
                         self.simulations[name].append(
-                            (posterior[name].sel(chain=0) < prior_draw[name]).sum(axis=0).values
+                            (posterior[name].sel(chain=0) < prior_draw[name])
+                            .sum(axis=0)
+                            .values
                         )
                 self._simulations_complete += 1
                 progress.update()
         finally:
             self.simulations = {
-                k: v[: self._simulations_complete] for k, v in self.simulations.items()
+                k: v[: self._simulations_complete]
+                for k, v in self.simulations.items()
             }
             progress.close()
 
@@ -143,11 +181,13 @@ def plot_results(self, kind="ecdf", var_names=None, color="C0"):
         simulations
             The SBC.simulations dictionary.
         kind : str
-            What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
+            What kind of plot to make. Supported values are 'ecdf' (default)
+            and 'hist'
         var_names : list[str]
             Variables to plot (defaults to all)
         figsize : tuple
-            Figure size for the plot. If None, it will be defined automatically.
+            Figure size for the plot. If None, it will be defined
+            automatically.
         color : str
             Color to use for the eCDF or histogram
 
@@ -156,4 +196,10 @@ def plot_results(self, kind="ecdf", var_names=None, color="C0"):
         fig, axes
             matplotlib figure and axes
         """
-        return plot_results(self.simulations, self.num_samples, kind=kind, var_names=var_names, color=color)
+        return plot_results(
+            self.simulations,
+            self.num_samples,
+            kind=kind,
+            var_names=var_names,
+            color=color,
+        )
diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index f7cb3ec..da708c8 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -1,4 +1,5 @@
 """Plots for the simulation based calibration"""
+
 import itertools
 
 import arviz as az
@@ -7,7 +8,9 @@
 from scipy.special import bdtrik
 
 
-def plot_results(simulations, ndraws, kind="ecdf", var_names=None, figsize=None, color="C0"):
+def plot_results(
+    simulations, ndraws, kind="ecdf", var_names=None, figsize=None, color="C0"
+):
     """Visual diagnostic for SBC.
 
     Currently it support two options: `ecdf` for the empirical CDF plots
@@ -22,7 +25,8 @@ def plot_results(simulations, ndraws, kind="ecdf", var_names=None, figsize=None,
     ndraws : int
         Number of draws in each posterior predictive sample
     kind : str
-        What kind of plot to make. Supported values are 'ecdf' (default) and 'hist'
+        What kind of plot to make. Supported values are 'ecdf' (default)
+        and 'hist'
     var_names : list[str]
         Variables to plot (defaults to all)
     figsize : tuple
@@ -53,13 +57,15 @@ def plot_results(simulations, ndraws, kind="ecdf", var_names=None, figsize=None,
 
     if n_plots > 1:
         if figsize is None:
-            figsize=(8, n_plots*1.0)
+            figsize = (8, n_plots * 1.0)
 
-        fig, axes = plt.subplots(nrows=(n_plots + 1) // 2, ncols=2, figsize=figsize, sharex=True)
+        fig, axes = plt.subplots(
+            nrows=(n_plots + 1) // 2, ncols=2, figsize=figsize, sharex=True
+        )
         axes = axes.flatten()
     else:
         if figsize is None:
-            figsize=(8, 1.5)
+            figsize = (8, 1.5)
 
         fig, axes = plt.subplots(nrows=1, ncols=1, figsize=figsize)
         axes = [axes]
@@ -69,28 +75,28 @@ def plot_results(simulations, ndraws, kind="ecdf", var_names=None, figsize=None,
 
     idx = 0
     for var_name, var_data in sims.items():
-        plot_idxs = list(itertools.product(*(np.arange(s) for s in var_data.shape[1:])))
-        if len(plot_idxs) > 1:
-            has_dims = True
-        else:
-            has_dims = False
+        plot_idxs = list(
+            itertools.product(*(np.arange(s) for s in var_data.shape[1:]))
+        )
 
         for indices in plot_idxs:
-            if has_dims:
-                dim_label = f'{var_name}[{"][".join(map(str, indices))}]'
+            if len(plot_idxs) > 1:  # has dims
+                dim_label = f"{var_name}[{']['.join(map(str, indices))}]"
             else:
                 dim_label = var_name
             ax = axes[idx]
             ary = var_data[(...,) + indices]
             if kind == "ecdf":
-                az.plot_ecdf(ary,
-                            cdf=cdf,
-                            difference = True,
-                            pit = True,
-                            confidence_bands = "auto",
-                            plot_kwargs={"color":color},
-                            fill_kwargs={"color":color},
-                            ax=ax)
+                az.plot_ecdf(
+                    ary,
+                    cdf=cdf,
+                    difference=True,
+                    pit=True,
+                    confidence_bands="auto",
+                    plot_kwargs={"color": color},
+                    fill_kwargs={"color": color},
+                    ax=ax,
+                )
             else:
                 hist(ary, color=color, ax=ax)
             ax.set_title(dim_label)
@@ -102,6 +108,7 @@ def plot_results(simulations, ndraws, kind="ecdf", var_names=None, figsize=None,
 
     return fig, axes
 
+
 def hist(ary, color, ax):
     hist, bins = np.histogram(ary, bins="auto")
     bin_centers = 0.5 * (bins[:-1] + bins[1:])
@@ -109,15 +116,29 @@ def hist(ary, color, ax):
     len_bins = len(bins)
     n_sims = len(ary)
 
-    band = np.ceil(bdtrik([0.025, 0.5, 0.975], n_sims, 1/len_bins))
-    ax.bar(bin_centers, hist, width=bins[1] - bins[0], color=color, edgecolor='black')
+    band = np.ceil(bdtrik([0.025, 0.5, 0.975], n_sims, 1 / len_bins))
+    ax.bar(
+        bin_centers,
+        hist,
+        width=bins[1] - bins[0],
+        color=color,
+        edgecolor="black",
+    )
     ax.axhline(band[1], color="0.5", ls="--")
-    ax.fill_between(np.linspace(0, max_rank, len_bins), band[0], band[2], color="0.5", alpha=0.5)
+    ax.fill_between(
+        np.linspace(0, max_rank, len_bins),
+        band[0],
+        band[2],
+        color="0.5",
+        alpha=0.5,
+    )
 
 
-class UniformCDF():
+class UniformCDF:
     def __init__(self, upper_bound):
         self.upper_bound = upper_bound
 
     def __call__(self, x):
-        return np.where(x < 0, 0, np.where(x > self.upper_bound, 1, x / self.upper_bound))
+        return np.where(
+            x < 0, 0, np.where(x > self.upper_bound, 1, x / self.upper_bound)
+        )
diff --git a/tests/test_sbc.py b/tests/test_sbc.py
index b7e690c..78ad49c 100644
--- a/tests/test_sbc.py
+++ b/tests/test_sbc.py
@@ -3,9 +3,9 @@
 
 ```math
 \begin{aligned}
-\mu &\sim \text{Normal}(0, 1), \\
-z &\sim \text{Normal}(\mu, 1).
-\end{aligned}
+\\mu &\\sim \text{Normal}(0, 1), \\
+z &\\sim \text{Normal}(\\mu, 1).
+\\end{aligned}
 ```
 """
 
@@ -20,21 +20,26 @@
 @pytest.fixture
 def simple_model():
     def model(y=None):
-        mu = numpyro.sample('mu', numpyro.distributions.Normal(0, 1))
-        numpyro.sample('z', numpyro.distributions.Normal(mu, 1), obs=y)
+        mu = numpyro.sample("mu", numpyro.distributions.Normal(0, 1))
+        numpyro.sample("z", numpyro.distributions.Normal(mu, 1), obs=y)
+
     return model
 
+
 @pytest.fixture
 def mcmc_kernel(simple_model):
     return NUTS(simple_model)
 
+
 @pytest.fixture
 def observed_vars():
-    return {"y" : "z"}
+    return {"y": "z"}
+
 
 @pytest.fixture
 def sbc_instance(mcmc_kernel, observed_vars):
-    return SBC(mcmc_kernel, y = None, observed_vars=observed_vars)
+    return SBC(mcmc_kernel, y=None, observed_vars=observed_vars)
+
 
 def test_sbc_initialization(sbc_instance, mcmc_kernel, observed_vars):
     """
@@ -43,9 +48,12 @@ def test_sbc_initialization(sbc_instance, mcmc_kernel, observed_vars):
     assert sbc_instance.mcmc_kernel == mcmc_kernel
     assert sbc_instance.observed_vars == observed_vars
     assert sbc_instance.num_simulations == 10
-    assert sbc_instance.sample_kwargs == dict(num_warmup=500, num_samples=100, progress_bar=False)
+    assert sbc_instance.sample_kwargs == dict(
+        num_warmup=500, num_samples=100, progress_bar=False
+    )
     assert sbc_instance._simulations_complete == 0
 
+
 def test_get_prior_predictive_samples(sbc_instance):
     """
     Test that the prior and prior predictive samples are generated correctly.
@@ -54,6 +62,7 @@ def test_get_prior_predictive_samples(sbc_instance):
     assert "y" in prior_pred
     assert "mu" in prior
 
+
 def test_get_posterior_samples(sbc_instance):
     """
     Test that the posterior samples are generated correctly.
@@ -64,6 +73,7 @@ def test_get_posterior_samples(sbc_instance):
     assert "posterior" in idata
     assert num_samples == 100
 
+
 def test_run_simulations(sbc_instance):
     """
     Test that the simulations for SBC are run correctly.
@@ -72,6 +82,7 @@ def test_run_simulations(sbc_instance):
     assert sbc_instance._simulations_complete == sbc_instance.num_simulations
     assert "mu" in sbc_instance.simulations
 
+
 def test_plot_results(sbc_instance):
     """
     Test that the results are plotted.

From a6b9e8815f318a649186123b10d59ed8940f9b76 Mon Sep 17 00:00:00 2001
From: AFg6K7h4fhy2 <127630341+AFg6K7h4fhy2@users.noreply.github.com>
Date: Tue, 11 Feb 2025 13:47:51 -0500
Subject: [PATCH 10/25] light docstring formatting

---
 forecasttools/sbc.py       | 9 +++++----
 forecasttools/sbc_plots.py | 4 ++--
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index da02fa3..ab8fbb1 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -19,8 +19,9 @@ def __init__(
         sample_kwargs=None,
         seed=1234,
         **kwargs,
-    ):
-        """Set up class for doing SBC.
+    ) -> None:
+        """
+        Set up class for doing SBC.
         Based on simulation based calibration (Talts et. al. 2018) in PyMC.
 
         Parameters
@@ -169,13 +170,13 @@ def run_simulations(self):
             progress.close()
 
     def plot_results(self, kind="ecdf", var_names=None, color="C0"):
-        """Visual diagnostic for SBC.
+        """
+        Visual diagnostic for SBC.
 
         Currently it support two options: `ecdf` for the empirical CDF plots
         of the difference between prior and posterior. `hist` for the rank
         histogram.
 
-
         Parameters
         ----------
         simulations
diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index da708c8..6cc748a 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -11,13 +11,13 @@
 def plot_results(
     simulations, ndraws, kind="ecdf", var_names=None, figsize=None, color="C0"
 ):
-    """Visual diagnostic for SBC.
+    """
+    Visual diagnostic for SBC.
 
     Currently it support two options: `ecdf` for the empirical CDF plots
     of the difference between prior and posterior. `hist` for the rank
     histogram.
 
-
     Parameters
     ----------
     simulations : dict[str] -> listlike

From 9e83506a5ebe0abfd9091c5fc52cd667c89a06e0 Mon Sep 17 00:00:00 2001
From: AFg6K7h4fhy2 <127630341+AFg6K7h4fhy2@users.noreply.github.com>
Date: Tue, 11 Feb 2025 13:57:08 -0500
Subject: [PATCH 11/25] revert then modify seed param edit

---
 forecasttools/sbc.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index ab8fbb1..ba3e388 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -17,7 +17,7 @@ def __init__(
         observed_vars: dict[str, str],
         num_simulations=10,
         sample_kwargs=None,
-        seed=1234,
+        seed=None,
         **kwargs,
     ) -> None:
         """
@@ -40,7 +40,7 @@ def __init__(
             Arguments passed to `numpyro.sample`. Defaults to
             `dict(num_warmup=500, num_samples=100, progress_bar = False)`.
             Which assumes a MCMC sampler e.g. NUTS.
-        seed : int
+        seed : random.PRNGKey
             Random seed.
         kwargs : dict
             Keyword arguments passed to `numpyro` models.
@@ -49,8 +49,8 @@ def __init__(
             sample_kwargs = dict(
                 num_warmup=500, num_samples=100, progress_bar=False
             )
-        seed = random.PRNGKey(seed)
-
+        if seed is None:
+            seed = random.PRNGKey(1234)
         self.mcmc_kernel = mcmc_kernel
         if not hasattr(mcmc_kernel, "model"):
             raise ValueError(

From e798c850c581ccf4e159d1903ffbd32350647026 Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Wed, 12 Feb 2025 09:54:54 +0000
Subject: [PATCH 12/25] Update forecasttools/sbc_plots.py

Co-authored-by: upx3 (CFA) <127630341+AFg6K7h4fhy2@users.noreply.github.com>
---
 forecasttools/sbc_plots.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index 6cc748a..67fecb5 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -1,4 +1,6 @@
-"""Plots for the simulation based calibration"""
+"""
+Plots for the simulation based calibration
+"""
 
 import itertools
 

From da3750bf2df64d60ac967df4af9f1457f5422e32 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 09:56:10 +0000
Subject: [PATCH 13/25] Update sbc_model_checking.qmd

---
 notebooks/sbc_model_checking.qmd | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/notebooks/sbc_model_checking.qmd b/notebooks/sbc_model_checking.qmd
index 4522595..2c1dd09 100644
--- a/notebooks/sbc_model_checking.qmd
+++ b/notebooks/sbc_model_checking.qmd
@@ -35,7 +35,7 @@ $$
 F^{(-1)}(\theta[k]) \sim \mathcal{U}(0,1).
 $$
 
-Where $F^{(-1)}$ is the (pseudo) inverse distribution function of the known distribution $p(\theta[k])$. This is often call ed the probability integral transform (PIT). Significant (in a classical statistics sense) deviations in the PIT of parameters from the uniform distribution form evidence that either the model is incorrectly coded, its priors are cover hard to sample regimes, and/or the Bayesian inference method is approximate/invalid. See Talts *et al* [@talts] for more details on understanding the returned tests.
+Where $F^{(-1)}$ is the (pseudo) inverse distribution function of the known distribution $p(\theta[k])$. This is often called the probability integral transform (PIT). Significant (in a classical statistics sense) deviations in the PIT of parameters from the uniform distribution form evidence that either the model is incorrectly coded, its priors are cover hard to sample regimes, and/or the Bayesian inference method is approximate/invalid. See Talts *et al* [@talts] for more details on understanding the returned tests.
 
 More precisely, in SBC:
 

From a1adf7bbb0fe099d2a4cc3cec985c70b0a00bec2 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:07:47 +0000
Subject: [PATCH 14/25] update about rank statistics

---
 forecasttools/sbc.py             | 10 ++++++----
 notebooks/sbc_model_checking.qmd |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index ba3e388..b0475fd 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -36,13 +36,13 @@ def __init__(
             Positional arguments passed to `numpyro.sample`.
         num_simulations : int
             How many simulations to run for SBC.
-        sample_kwargs : dict[str] -> Any
+        sample_kwargs : dict[str, Any]
             Arguments passed to `numpyro.sample`. Defaults to
             `dict(num_warmup=500, num_samples=100, progress_bar = False)`.
             Which assumes a MCMC sampler e.g. NUTS.
         seed : random.PRNGKey
             Random seed.
-        kwargs : dict
+        kwargs : dict[str, Any]
             Keyword arguments passed to `numpyro` models.
         """
         if sample_kwargs is None:
@@ -149,17 +149,19 @@ def run_simulations(self):
                 for name in prior:
                     num_dims = jnp.ndim(prior_draw[name])
                     if num_dims == 0:
-                        self.simulations[name].append(
+                        rank_statistics = (
                             (posterior[name].sel(chain=0) < prior_draw[name])
                             .sum()
                             .values
                         )
+                        self.simulations[name].append(rank_statistics)
                     else:
-                        self.simulations[name].append(
+                        rank_statistics = (
                             (posterior[name].sel(chain=0) < prior_draw[name])
                             .sum(axis=0)
                             .values
                         )
+                        self.simulations[name].append(rank_statistics)
                 self._simulations_complete += 1
                 progress.update()
         finally:
diff --git a/notebooks/sbc_model_checking.qmd b/notebooks/sbc_model_checking.qmd
index 2c1dd09..3e133a4 100644
--- a/notebooks/sbc_model_checking.qmd
+++ b/notebooks/sbc_model_checking.qmd
@@ -44,7 +44,7 @@ More precisely, in SBC:
 3.  Generate the probability integral transform (PIT) of each of $k = 1,\dots,P$ parameters $\theta[k]$ with respect to the known prior distribution.
 4.  Assess the PIT distributions for deviation against an assumption of being $\mathcal{U}([0,1])$.
 
-In `SBC` we from the PIT by looking at the distribution of proportion of posterior samples that are less than the "true" parameter values cached along with the generated data $y^{(i)}$: $P(\theta_p^{(i)}[k] < \theta^{(i)}[k])$. This is a more convenient approach in `numpyro` than trying to solve the inverse distribution function directly.
+In `SBC` we form the PIT by looking at the distribution of proportion of posterior samples that are less than the "true" parameter values cached along with the generated data $y^{(i)}$: $P(\theta_p^{(i)}[k] < \theta^{(i)}[k])$. These are commonly called the rank statistics of the sampling process. Using rank statistics is a more convenient approach in `numpyro` than trying to solve the inverse distribution function directly.
 
 ## Example: Eight schools
 

From cd242ea3f557f3ae0960a82493d7a2c991a44c67 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:10:41 +0000
Subject: [PATCH 15/25] type-hint change

---
 forecasttools/sbc_plots.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index 67fecb5..6b0596b 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -22,7 +22,7 @@ def plot_results(
 
     Parameters
     ----------
-    simulations : dict[str] -> listlike
+    simulations : dict[str, Any]
         The SBC.simulations dictionary.
     ndraws : int
         Number of draws in each posterior predictive sample

From 241f31cc991d37a9a189e70a930c842941e615db Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:14:18 +0000
Subject: [PATCH 16/25] change axes iterator label

---
 forecasttools/sbc_plots.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index 6b0596b..443d56a 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -105,8 +105,8 @@ def plot_results(
             ax.set_yticks([])
             idx += 1
 
-    for j in range(n_plots, len(axes)):
-        fig.delaxes(axes[j])
+    for extra_ax in range(n_plots, len(axes)):
+        fig.delaxes(axes[extra_ax])
 
     return fig, axes
 

From ef77270b5c59b0295d4be733b6ed81a04534c7bd Mon Sep 17 00:00:00 2001
From: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Date: Wed, 12 Feb 2025 11:26:38 +0000
Subject: [PATCH 17/25] Update forecasttools/sbc_plots.py

Co-authored-by: upx3 (CFA) <127630341+AFg6K7h4fhy2@users.noreply.github.com>
---
 forecasttools/sbc_plots.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index 443d56a..12bef0b 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -138,6 +138,8 @@ def hist(ary, color, ax):
 
 class UniformCDF:
     def __init__(self, upper_bound):
+        if upper_bound <= 0:
+            raise ValueError(f"Upper bound must be positive; got {upper_bound}.")
         self.upper_bound = upper_bound
 
     def __call__(self, x):

From fad65397b4f8604a14898759c196d15a169c2cde Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:33:02 +0000
Subject: [PATCH 18/25] docstrings for UniformCDF

---
 forecasttools/sbc_plots.py | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index 12bef0b..dd5a8f4 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -137,12 +137,46 @@ def hist(ary, color, ax):
 
 
 class UniformCDF:
+    """
+    A class to represent the cumulative distribution function (CDF) of a
+    uniform distribution.
+    """
+
     def __init__(self, upper_bound):
+        """
+        Constructs all the necessary attributes for the UniformCDF object.
+
+        Parameters
+        ----------
+        upper_bound : float
+            The upper bound of the uniform distribution. Must be positive.
+
+        Raises
+        ------
+        ValueError
+            If upper_bound is not positive.
+        """
         if upper_bound <= 0:
-            raise ValueError(f"Upper bound must be positive; got {upper_bound}.")
+            raise ValueError(
+                f"Upper bound must be positive; got {upper_bound}."
+            )
         self.upper_bound = upper_bound
 
     def __call__(self, x):
+        """
+        Evaluates the CDF at a given value x.
+
+        Parameters
+        ----------
+        x : array-like
+            The value(s) at which to evaluate the CDF.
+
+        Returns
+        -------
+        array-like
+            The CDF evaluated at x. Returns 0 if x < 0, 1 if x > upper_bound,
+            and x / upper_bound otherwise.
+        """
         return np.where(
             x < 0, 0, np.where(x > self.upper_bound, 1, x / self.upper_bound)
         )

From 6da1696f70a3e387bd8cf1335b0ce7e20e57e28a Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:38:54 +0000
Subject: [PATCH 19/25] typehint for _get_prior_predictive_samples return

---
 forecasttools/sbc.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index b0475fd..99581c8 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -79,7 +79,9 @@ def __init__(
         self._sampler_rng = sampler_rng
         self.num_samples = None
 
-    def _get_prior_predictive_samples(self):
+    def _get_prior_predictive_samples(
+        self,
+    ) -> tuple[dict[str, any], dict[str, any]]:
         """
         Generate samples to use for the simulations by prior predictive
         sampling. Then splits between observed and unobserved variables based

From d2666705b7f0389c1a74f342a2dda7c5c9e9e313 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:43:32 +0000
Subject: [PATCH 20/25] type hints for _get_posterior_samples

---
 forecasttools/sbc.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 99581c8..3fe1288 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -86,6 +86,11 @@ def _get_prior_predictive_samples(
         Generate samples to use for the simulations by prior predictive
         sampling. Then splits between observed and unobserved variables based
         on the `observed_vars` attribute.
+
+        Returns
+        -------
+        tuple[dict[str, any], dict[str, any]]
+            The prior and prior predictive samples.
         """
         prior_predictive_fn = numpyro.infer.Predictive(
             self.mcmc_kernel.model, num_samples=self.num_simulations
@@ -103,12 +108,27 @@ def _get_prior_predictive_samples(
         }
         return prior, prior_pred
 
-    def _get_posterior_samples(self, seed, prior_predictive_draw):
+    def _get_posterior_samples(
+        self, seed: random.PRNGKey, prior_predictive_draw: dict[str, any]
+    ) -> tuple[az.InferenceData, int]:
         """
         Generate posterior samples conditioned to a prior predictive sample.
         This returns the posterior samples and the number of samples. The
         number of samples are used in scaling plotting and checking that each
         inference draw has the same number of samples.
+
+        Parameters
+        ----------
+        seed : random.PRNGKey
+            Random seed for MCMC sampling.
+        prior_predictive_draw : dict[str, any]
+            Prior predictive samples.
+
+        Returns
+        -------
+        tuple[az.InferenceData, int]
+            Posterior samples as an arviz InferenceData object, with the count
+            of posterior samples.
         """
         mcmc = MCMC(self.mcmc_kernel, **self.sample_kwargs)
         obs_vars = {**self.kwargs, **prior_predictive_draw}

From 501ba7cfc787855d78cb24a64ea997fff337858e Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Wed, 12 Feb 2025 11:45:15 +0000
Subject: [PATCH 21/25] typehint for run_simulations

---
 forecasttools/sbc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 3fe1288..2e99664 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -137,7 +137,7 @@ def _get_posterior_samples(
         idata = az.from_numpyro(mcmc)
         return idata, num_samples
 
-    def run_simulations(self):
+    def run_simulations(self) -> None:
         """
         The main method of `SBC` class that runs the simulations for
         simulation based calibration and fills the `simulations` attribute

From 8346deb01810f816d97e56acc18fbe10c1412b6f Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 14 Feb 2025 17:55:19 +0000
Subject: [PATCH 22/25] refactor SBC inner loop

---
 forecasttools/sbc.py | 54 +++++++++++++++++++++++---------------------
 tests/test_sbc.py    |  3 +--
 2 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 2e99664..063cd5b 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -134,8 +134,33 @@ def _get_posterior_samples(
         obs_vars = {**self.kwargs, **prior_predictive_draw}
         mcmc.run(seed, *self.args, **obs_vars)
         num_samples = mcmc.num_samples
+        # Check that the number of samples is consistent
+        if self.num_samples is None:
+            self.num_samples = num_samples
+        if self.num_samples != num_samples:
+            raise ValueError(
+                "The number of samples from the posterior is not consistent."
+            )
         idata = az.from_numpyro(mcmc)
-        return idata, num_samples
+        return idata
+
+    def increment_rank_statistics(self, prior_draw, posterior):
+        for name in prior_draw:
+            num_dims = jnp.ndim(prior_draw[name])
+            if num_dims == 0:
+                rank_statistics = (
+                    (posterior[name].sel(chain=0) < prior_draw[name])
+                    .sum()
+                    .values
+                )
+                self.simulations[name].append(rank_statistics)
+            else:
+                rank_statistics = (
+                    (posterior[name].sel(chain=0) < prior_draw[name])
+                    .sum(axis=0)
+                    .values
+                )
+                self.simulations[name].append(rank_statistics)
 
     def run_simulations(self) -> None:
         """
@@ -157,33 +182,10 @@ def run_simulations(self) -> None:
                 prior_predictive_draw = {
                     k: v[idx] for k, v in prior_pred.items()
                 }
-                idata, num_samples = self._get_posterior_samples(
+                idata = self._get_posterior_samples(
                     sampler_seeds[idx], prior_predictive_draw
                 )
-                if self.num_samples is None:
-                    self.num_samples = num_samples
-                if self.num_samples != num_samples:
-                    raise ValueError(
-                        "The number of samples from the posterior is not"
-                        " consistent."
-                    )
-                posterior = idata["posterior"]
-                for name in prior:
-                    num_dims = jnp.ndim(prior_draw[name])
-                    if num_dims == 0:
-                        rank_statistics = (
-                            (posterior[name].sel(chain=0) < prior_draw[name])
-                            .sum()
-                            .values
-                        )
-                        self.simulations[name].append(rank_statistics)
-                    else:
-                        rank_statistics = (
-                            (posterior[name].sel(chain=0) < prior_draw[name])
-                            .sum(axis=0)
-                            .values
-                        )
-                        self.simulations[name].append(rank_statistics)
+                self.increment_rank_statistics(prior_draw, idata["posterior"])
                 self._simulations_complete += 1
                 progress.update()
         finally:
diff --git a/tests/test_sbc.py b/tests/test_sbc.py
index 78ad49c..13824eb 100644
--- a/tests/test_sbc.py
+++ b/tests/test_sbc.py
@@ -69,9 +69,8 @@ def test_get_posterior_samples(sbc_instance):
     """
     prior, prior_pred = sbc_instance._get_prior_predictive_samples()
     seed = random.PRNGKey(0)
-    idata, num_samples = sbc_instance._get_posterior_samples(seed, prior_pred)
+    idata = sbc_instance._get_posterior_samples(seed, prior_pred)
     assert "posterior" in idata
-    assert num_samples == 100
 
 
 def test_run_simulations(sbc_instance):

From 9efc9bb3af7f214be32a2b3e23b24dd5e19b1f6f Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 14 Feb 2025 18:57:31 +0000
Subject: [PATCH 23/25] Add unit test and docstring

---
 forecasttools/sbc.py | 14 ++++++++++++--
 tests/test_sbc.py    | 16 ++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 063cd5b..80b9a87 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -144,7 +144,17 @@ def _get_posterior_samples(
         idata = az.from_numpyro(mcmc)
         return idata
 
-    def increment_rank_statistics(self, prior_draw, posterior):
+    def _increment_rank_statistics(self, prior_draw, posterior) -> None:
+        """
+        Increment the rank statistics for each parameter in the prior draw.
+
+        This method updates the `self.simulations` dictionary with the rank
+        statistics for each parameter in the `prior_draw` compared to the
+        `posterior`.
+
+        Returns:
+            None
+        """
         for name in prior_draw:
             num_dims = jnp.ndim(prior_draw[name])
             if num_dims == 0:
@@ -185,7 +195,7 @@ def run_simulations(self) -> None:
                 idata = self._get_posterior_samples(
                     sampler_seeds[idx], prior_predictive_draw
                 )
-                self.increment_rank_statistics(prior_draw, idata["posterior"])
+                self._increment_rank_statistics(prior_draw, idata["posterior"])
                 self._simulations_complete += 1
                 progress.update()
         finally:
diff --git a/tests/test_sbc.py b/tests/test_sbc.py
index 13824eb..e91d7b5 100644
--- a/tests/test_sbc.py
+++ b/tests/test_sbc.py
@@ -73,6 +73,22 @@ def test_get_posterior_samples(sbc_instance):
     assert "posterior" in idata
 
 
+def test_increment_rank_statistics(sbc_instance):
+    """
+    Test that the rank statistics are incremented correctly.
+    """
+    prior, prior_pred = sbc_instance._get_prior_predictive_samples()
+    seed = random.PRNGKey(0)
+    idata = sbc_instance._get_posterior_samples(seed, prior_pred)
+    sbc_instance.simulations = {name: [] for name in prior}
+    prior_draw = {k: v[0] for k, v in prior.items()}
+    sbc_instance._increment_rank_statistics(prior_draw, idata["posterior"])
+
+    for name in prior_draw:
+        assert name in sbc_instance.simulations
+        assert len(sbc_instance.simulations[name]) == 1
+
+
 def test_run_simulations(sbc_instance):
     """
     Test that the simulations for SBC are run correctly.

From 250e49dfa61b8bbc2fd036d7e2266e1ba9be7013 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 14 Feb 2025 19:12:46 +0000
Subject: [PATCH 24/25] add inspection mode

---
 forecasttools/sbc.py | 13 ++++++++++++-
 tests/test_sbc.py    | 17 +++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/forecasttools/sbc.py b/forecasttools/sbc.py
index 80b9a87..1e9d4e1 100644
--- a/forecasttools/sbc.py
+++ b/forecasttools/sbc.py
@@ -18,6 +18,7 @@ def __init__(
         num_simulations=10,
         sample_kwargs=None,
         seed=None,
+        inspection_mode=False,
         **kwargs,
     ) -> None:
         """
@@ -71,13 +72,18 @@ def __init__(
 
         self.num_simulations = num_simulations
         self.sample_kwargs = sample_kwargs
-
+        # Initialize the simulations and random seeds
         self.simulations = {}
         self._simulations_complete = 0
         prior_pred_rng, sampler_rng = random.split(seed)
         self._prior_pred_rng = prior_pred_rng
         self._sampler_rng = sampler_rng
         self.num_samples = None
+        # Set the inspection mode
+        # if in inspection mode, store all idata objects from fitting
+        self.inspection_mode = inspection_mode
+        if inspection_mode:
+            self.idatas = []
 
     def _get_prior_predictive_samples(
         self,
@@ -185,6 +191,9 @@ def run_simulations(self) -> None:
             initial=self._simulations_complete,
             total=self.num_simulations,
         )
+        if self.inspection_mode:
+            self.prior = prior
+            self.prior_pred = prior_pred
         try:
             while self._simulations_complete < self.num_simulations:
                 idx = self._simulations_complete
@@ -195,6 +204,8 @@ def run_simulations(self) -> None:
                 idata = self._get_posterior_samples(
                     sampler_seeds[idx], prior_predictive_draw
                 )
+                if self.inspection_mode:
+                    self.idatas.append(idata)
                 self._increment_rank_statistics(prior_draw, idata["posterior"])
                 self._simulations_complete += 1
                 progress.update()
diff --git a/tests/test_sbc.py b/tests/test_sbc.py
index e91d7b5..d5831ed 100644
--- a/tests/test_sbc.py
+++ b/tests/test_sbc.py
@@ -41,6 +41,13 @@ def sbc_instance(mcmc_kernel, observed_vars):
     return SBC(mcmc_kernel, y=None, observed_vars=observed_vars)
 
 
+@pytest.fixture
+def sbc_instance_inspection_on(mcmc_kernel, observed_vars):
+    return SBC(
+        mcmc_kernel, y=None, observed_vars=observed_vars, inspection_mode=True
+    )
+
+
 def test_sbc_initialization(sbc_instance, mcmc_kernel, observed_vars):
     """
     Test that the SBC class is initialized correctly.
@@ -98,6 +105,16 @@ def test_run_simulations(sbc_instance):
     assert "mu" in sbc_instance.simulations
 
 
+def test_run_simulations_with_inspection(sbc_instance_inspection_on):
+    """
+    Test that the simulations for SBC are run correctly.
+    """
+    sbc_instance_inspection_on.run_simulations()
+    assert isinstance(sbc_instance_inspection_on.idatas, list)
+    assert isinstance(sbc_instance_inspection_on.prior, dict)
+    assert isinstance(sbc_instance_inspection_on.prior_pred, dict)
+
+
 def test_plot_results(sbc_instance):
     """
     Test that the results are plotted.

From 7ee709ca60307462e44a4419b92eb447bb234e1d Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 14 Feb 2025 19:29:39 +0000
Subject: [PATCH 25/25] remove hand rolled uniformCDF class

Replaced with numpyro.distributions.DiscreteUniform
---
 forecasttools/sbc_plots.py | 49 ++------------------------------------
 1 file changed, 2 insertions(+), 47 deletions(-)

diff --git a/forecasttools/sbc_plots.py b/forecasttools/sbc_plots.py
index dd5a8f4..43f68e2 100644
--- a/forecasttools/sbc_plots.py
+++ b/forecasttools/sbc_plots.py
@@ -7,6 +7,7 @@
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
+import numpyro.distributions as dist
 from scipy.special import bdtrik
 
 
@@ -73,7 +74,7 @@ def plot_results(
         axes = [axes]
 
     if kind == "ecdf":
-        cdf = UniformCDF(ndraws)
+        cdf = dist.DiscreteUniform(high=ndraws).cdf
 
     idx = 0
     for var_name, var_data in sims.items():
@@ -134,49 +135,3 @@ def hist(ary, color, ax):
         color="0.5",
         alpha=0.5,
     )
-
-
-class UniformCDF:
-    """
-    A class to represent the cumulative distribution function (CDF) of a
-    uniform distribution.
-    """
-
-    def __init__(self, upper_bound):
-        """
-        Constructs all the necessary attributes for the UniformCDF object.
-
-        Parameters
-        ----------
-        upper_bound : float
-            The upper bound of the uniform distribution. Must be positive.
-
-        Raises
-        ------
-        ValueError
-            If upper_bound is not positive.
-        """
-        if upper_bound <= 0:
-            raise ValueError(
-                f"Upper bound must be positive; got {upper_bound}."
-            )
-        self.upper_bound = upper_bound
-
-    def __call__(self, x):
-        """
-        Evaluates the CDF at a given value x.
-
-        Parameters
-        ----------
-        x : array-like
-            The value(s) at which to evaluate the CDF.
-
-        Returns
-        -------
-        array-like
-            The CDF evaluated at x. Returns 0 if x < 0, 1 if x > upper_bound,
-            and x / upper_bound otherwise.
-        """
-        return np.where(
-            x < 0, 0, np.where(x > self.upper_bound, 1, x / self.upper_bound)
-        )