From 4541170e76e5e624bc20417d852b0f6e06ab8a46 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Sun, 26 Jul 2026 23:05:11 +0200
Subject: [PATCH 01/16] refactor: ensure numpyro stores posteriors and priors
 as list[xr.Dataset] and xr.Dataset like pymc

---
 simuk/numpyro_adapter.py | 26 ++++++++++++++++----------
 simuk/sbc.py             |  3 ++-
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/simuk/numpyro_adapter.py b/simuk/numpyro_adapter.py
index 20ca7a5..160f8d4 100644
--- a/simuk/numpyro_adapter.py
+++ b/simuk/numpyro_adapter.py
@@ -4,7 +4,8 @@
 
 import jax
 import numpy as np
-from arviz_base import from_numpyro
+import xarray as xr
+from arviz_base import dict_to_dataset, extract, from_numpyro
 from numpyro.handlers import seed, trace
 from numpyro.infer import MCMC, Predictive
 
@@ -24,13 +25,14 @@ def __init__(self, data_dir, numpyro_model, model, simulator, single_seed):
     def compute_single_rank(self, transform, name, posterior, simulation_idx, ref_params):
         transformed_posterior = np.array(
             [
-                transform(name, posterior[name].sel(chain=0).isel(draw=i).values)
-                for i in range(posterior[name].sizes["draw"])
+                transform(name, posterior[name].isel(sample=i).values)
+                for i in range(posterior[name].sizes["sample"])
             ]
         )
-        return (transformed_posterior < transform(name, ref_params[name][simulation_idx])).sum(
-            axis=0
-        )
+        return (
+            transformed_posterior
+            < transform(name, ref_params[name].isel(sample=simulation_idx).values)
+        ).sum(axis=0)
 
     def get_posterior_predictive_samples(self, num_simulations, seeds, progress_bar):
         raise NotImplementedError("Posterior SBC is not implemented for numpyro")
@@ -51,9 +53,13 @@ def get_prior_predictive_samples(self, num_samples, seeds):
                 params = dict(zip(prior.keys(), vals))
                 params["seed"] = seeds[i]
                 results.append(self.simulator(**params))
-            prior_pred = {key: [result[key] for result in results] for key in results[0]}
+            prior_pred = {key: np.asarray([result[key] for result in results]) for key in results[0]}
         else:
             prior_pred = {k: v for k, v in samples.items() if k in self.observed_model_vars}
+
+        prior = dict_to_dataset(prior, sample_dims=["sample"])
+        prior_pred = dict_to_dataset(prior_pred, sample_dims=["sample"])
+
         return prior, prior_pred
 
     def _extract_model_info(self, single_seed):
@@ -126,14 +132,14 @@ def get_posterior_samples(
             if k in simulation_parameters.observed_model_vars
         }
         mcmc.run(rng_seed, **free_vars_data, **prior_predictive_args)
-        return from_numpyro(mcmc)["posterior"]
+        return extract(from_numpyro(mcmc), group="posterior", keep_dataset=True)
 
     def subsample(self, ref_params, predictive, seed, size):
         log.info("Slicing isn't implemented for numpyro, skipping it.")
         return ref_params, predictive
 
     def replicate(self, predictive, idx, simulation_params):
-        return {k: v[idx] for k, v in predictive.items()}
+        return {k: v.isel(sample=idx).values for k, v in predictive.items()}
 
     def stop_if_cant_run_without_simulator(self):
         if not self.observed_model_vars:
@@ -154,4 +160,4 @@ class NumpyroSimulationParams(NamedTuple):
     observed_vars: list[str]
     observed_model_vars: list[str]
     var_names: list[str]
-    ref_params: dict[str, jax.Array]
+    ref_params: xr.Dataset
diff --git a/simuk/sbc.py b/simuk/sbc.py
index 985c8c3..da7e89d 100644
--- a/simuk/sbc.py
+++ b/simuk/sbc.py
@@ -30,6 +30,7 @@
     pass
 
 import numpy as np
+import xarray as xr
 from arviz_base import from_dict
 from tqdm import tqdm
 
@@ -266,7 +267,7 @@ def __init__(
         self.sample_kwargs = sample_kwargs
         self.simulations = {name: [] for name in self.adapter.var_names}
         self._simulations_complete = 0
-        self.posteriors = []
+        self.posteriors: list[xr.Dataset] = []
         self.keep_fits = keep_fits
 
         if simulator is not None and not callable(simulator):

From d5ef04a4c375d7bc7299d53b56573b0501b64594 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Sun, 26 Jul 2026 23:10:59 +0200
Subject: [PATCH 02/16] chore: linting errors

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

diff --git a/simuk/numpyro_adapter.py b/simuk/numpyro_adapter.py
index 160f8d4..2cc4eee 100644
--- a/simuk/numpyro_adapter.py
+++ b/simuk/numpyro_adapter.py
@@ -53,7 +53,9 @@ def get_prior_predictive_samples(self, num_samples, seeds):
                 params = dict(zip(prior.keys(), vals))
                 params["seed"] = seeds[i]
                 results.append(self.simulator(**params))
-            prior_pred = {key: np.asarray([result[key] for result in results]) for key in results[0]}
+            prior_pred = {
+                key: np.asarray([result[key] for result in results]) for key in results[0]
+            }
         else:
             prior_pred = {k: v for k, v in samples.items() if k in self.observed_model_vars}
 

From 5caefba24c32ef4be73c91bb0d768aad3e29be5d Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Mon, 20 Jul 2026 17:22:58 +0200
Subject: [PATCH 03/16] feat: parameter recover plot

---
 simuk/__init__.py        |   1 +
 simuk/plots.py           | 325 +++++++++++++++++++++++++++++++++++++++
 simuk/tests/test_plot.py | 232 ++++++++++++++++++++++++++++
 3 files changed, 558 insertions(+)
 create mode 100644 simuk/plots.py
 create mode 100644 simuk/tests/test_plot.py

diff --git a/simuk/__init__.py b/simuk/__init__.py
index 9433981..8225be3 100644
--- a/simuk/__init__.py
+++ b/simuk/__init__.py
@@ -4,6 +4,7 @@
 Simulation based calibration and other tools for evaluation using synthetic data.
 """
 
+from simuk.plots import plot_parameter_recovery, plot_ecdf
 from simuk.sbc import SBC
 
 __version__ = "0.3.0"
diff --git a/simuk/plots.py b/simuk/plots.py
new file mode 100644
index 0000000..f06b43e
--- /dev/null
+++ b/simuk/plots.py
@@ -0,0 +1,325 @@
+"""Parameter recovery plotting for simulation-based calibration."""
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+from arviz_plots import plot_ecdf_pit
+
+
+def _extract_draws(posterior, name):
+    """Extract posterior draws for a parameter as a numpy array.
+
+    Handles both PyMC-style posteriors (with a ``sample`` dimension from
+    ``arviz_base.extract``) and NumPyro-style posteriors (with ``chain``
+    and ``draw`` dimensions from ``arviz_base.from_numpyro``).
+
+    Parameters
+    ----------
+    posterior : xarray.Dataset
+        Posterior draws for a single SBC iteration.
+    name : str
+        Parameter name.
+
+    Returns
+    -------
+    numpy.ndarray
+        Array with shape ``(n_draws, *param_shape)``.
+    """
+    var = posterior[name]
+    if "sample" in var.dims:
+        return np.asarray(var.values)
+    else:
+        stacked = var.stack(__sample__=("chain", "draw"))
+        return np.asarray(stacked.transpose("__sample__", ...).values)
+
+
+def _extract_true_value(ref_params, name, idx):
+    """Extract the true (reference) parameter value for a simulation.
+
+    Parameters
+    ----------
+    ref_params : xarray.Dataset or dict
+        Reference parameters. An xarray Dataset (PyMC) or a plain dict of
+        arrays (NumPyro).
+    name : str
+        Parameter name.
+    idx : int
+        Simulation index.
+
+    Returns
+    -------
+    numpy.ndarray
+        The true parameter value (0-d for scalars, 1-d+ for vectors).
+    """
+    if hasattr(ref_params, "data_vars"):
+        return np.asarray(ref_params[name].isel(sample=idx).values)
+    else:
+        return np.asarray(ref_params[name][idx])
+
+
+def plot_parameter_recovery(
+    sbc,
+    ci_prob=0.94,
+    point_estimate="mean",
+    transform=None,
+    var_names=None,
+    figsize=None,
+    axes=None,
+    if_show=True,
+):
+    r"""Plot parameter recovery: posterior point estimate vs true simulated value.
+
+    For each SBC iteration the true (reference) parameter value is plotted
+    on the x-axis and the posterior point estimate (mean or median) on the
+    y-axis, with credible-interval error bars.  Points whose credible
+    interval covers the true value are coloured blue; misses are red.
+
+    A *y = x* diagonal reference line indicates perfect recovery.
+
+    Parameters
+    ----------
+    sbc : simuk.SBC
+        An SBC instance that has completed simulations with
+        ``keep_fits=True``.
+    ci_prob : float, default 0.94
+        Probability mass of the credible interval (e.g. 0.94 for a 94 %
+        CI).  The interval is computed from posterior quantiles — no
+        distributional assumptions (e.g. normality) are made.
+    point_estimate : {"mean", "median"}, default "mean"
+        Which posterior point estimate to plot on the y-axis.
+    transform : callable, optional
+        ``(param_name, param_value) -> transformed_value``.  Applied to
+        both true values and posterior draws before summarising.  Defaults
+        to the transform set at SBC initialisation.
+    var_names : list[str], optional
+        Subset of parameter names to plot.  Defaults to all free
+        parameters.
+    figsize : tuple[float, float], optional
+        Figure size in inches.
+    axes : matplotlib Axes or array of Axes, optional
+        Pre-existing axes to draw on (one per scalar subplot).
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        The figure containing the parameter recovery subplot(s).
+
+    Notes
+    -----
+    **Credible intervals** are quantile-based: for a 94 % CI the lower
+    and upper bounds are the 3rd and 97th percentiles of the posterior
+    draws.
+
+    **Coverage** reported in subplot titles is the empirical fraction of
+    simulations whose CI contains the true value.  The "expected" range
+    uses the binomial standard error
+
+    .. math::
+
+        \mathrm{SE} = \sqrt{\frac{p\,(1-p)}{N}}
+
+    which accounts only for the sampling noise from having a finite number
+    *N* of SBC iterations.  This is a **lower bound** on the true
+    uncertainty: it ignores the additional noise introduced by estimating
+    each CI from a finite number of MCMC draws.  When the number of
+    posterior draws per fit is small, the actual uncertainty in empirical
+    coverage will be larger than the displayed ± range.
+    """
+    if not if_show:
+        matplotlib.use("Agg")
+    else:
+        matplotlib.use("TkAgg")
+
+    # ---- validation --------------------------------------------------------
+    if not sbc.keep_fits:
+        raise ValueError(
+            "`plot_parameter_recovery` requires the SBC instance to have "
+            "been run with `keep_fits=True`."
+        )
+    if not sbc.posteriors:
+        raise ValueError("No posteriors found. Call `sbc.run_simulations()` before plotting.")
+    if point_estimate not in ("mean", "median"):
+        raise ValueError(f"`point_estimate` must be 'mean' or 'median', got '{point_estimate}'")
+    if transform is None:
+        transform = sbc._transform
+    elif not callable(transform):
+        raise ValueError("`transform` should be a function or None")
+
+    ref_params = sbc.kept_simulation_params.ref_params
+    if var_names is None:
+        var_names = list(sbc.kept_simulation_params.var_names)
+
+    n_sims = len(sbc.posteriors)
+    alpha_lo = (1 - ci_prob) / 2
+    alpha_hi = 1 - alpha_lo
+
+    # ---- compute summaries -------------------------------------------------
+    summaries = {}
+    for name in var_names:
+        true_list, est_list, lo_list, hi_list = [], [], [], []
+
+        for idx in range(n_sims):
+            posterior = sbc.posteriors[idx]
+            draws = _extract_draws(posterior, name)
+            transformed_draws = np.array([transform(name, d) for d in draws])
+
+            if point_estimate == "mean":
+                est = np.mean(transformed_draws, axis=0)
+            else:
+                est = np.median(transformed_draws, axis=0)
+
+            ci_lo = np.quantile(transformed_draws, alpha_lo, axis=0)
+            ci_hi = np.quantile(transformed_draws, alpha_hi, axis=0)
+
+            true_val = _extract_true_value(ref_params, name, idx)
+            transformed_true = np.asarray(transform(name, true_val), dtype=float)
+
+            true_list.append(np.atleast_1d(transformed_true))
+            est_list.append(np.atleast_1d(est))
+            lo_list.append(np.atleast_1d(ci_lo))
+            hi_list.append(np.atleast_1d(ci_hi))
+
+        summaries[name] = {
+            "true_values": np.array(true_list),
+            "point_estimates": np.array(est_list),
+            "ci_lower": np.array(lo_list),
+            "ci_upper": np.array(hi_list),
+        }
+
+    # ---- determine subplot layout ------------------------------------------
+    subplot_specs = []  # (param_name, column_index, display_label)
+    for name in var_names:
+        n_elements = summaries[name]["true_values"].shape[1]
+        if n_elements == 1:
+            subplot_specs.append((name, 0, name))
+        else:
+            for j in range(n_elements):
+                subplot_specs.append((name, j, f"{name}[{j}]"))
+
+    n_subplots = len(subplot_specs)
+
+    # ---- create or validate axes -------------------------------------------
+    created_figure = axes is None
+    if created_figure:
+        n_cols = min(3, n_subplots)
+        n_rows = -(-n_subplots // n_cols)  # ceil division
+        if figsize is None:
+            figsize = (5 * n_cols, 4 * n_rows)
+        fig, axes_grid = plt.subplots(n_rows, n_cols, figsize=figsize, squeeze=False)
+        axes_flat = axes_grid.flatten()
+        for k in range(n_subplots, len(axes_flat)):
+            axes_flat[k].set_visible(False)
+    else:
+        axes_flat = np.atleast_1d(axes).flatten()
+        if len(axes_flat) < n_subplots:
+            raise ValueError(f"Need {n_subplots} axes but only {len(axes_flat)} provided.")
+        fig = axes_flat[0].get_figure()
+
+    # ---- plot each subplot -------------------------------------------------
+    binomial_se = np.sqrt(ci_prob * (1 - ci_prob) / n_sims)
+
+    for i, (name, col, label) in enumerate(subplot_specs):
+        ax = axes_flat[i]
+        s = summaries[name]
+
+        true_vals = s["true_values"][:, col]
+        estimates = s["point_estimates"][:, col]
+        ci_lo = s["ci_lower"][:, col]
+        ci_hi = s["ci_upper"][:, col]
+
+        covered = (ci_lo <= true_vals) & (true_vals <= ci_hi)
+        empirical_cov = covered.mean()
+
+        yerr = np.array(
+            [
+                np.maximum(estimates - ci_lo, 0),
+                np.maximum(ci_hi - estimates, 0),
+            ]
+        )
+
+        # Plot covered and missed points separately for the legend
+        mask_cov = covered
+        mask_miss = ~covered
+
+        if mask_cov.any():
+            ax.errorbar(
+                true_vals[mask_cov],
+                estimates[mask_cov],
+                yerr=yerr[:, mask_cov],
+                fmt="o",
+                color="C0",
+                markersize=4,
+                linewidth=1,
+                alpha=0.7,
+                capsize=2,
+                label="Covered",
+            )
+        if mask_miss.any():
+            ax.errorbar(
+                true_vals[mask_miss],
+                estimates[mask_miss],
+                yerr=yerr[:, mask_miss],
+                fmt="o",
+                color="C3",
+                markersize=4,
+                linewidth=1,
+                alpha=0.7,
+                capsize=2,
+                label="Missed",
+            )
+
+        # y = x reference line
+        all_vals = np.concatenate([true_vals, estimates, ci_lo, ci_hi])
+        vmin, vmax = all_vals.min(), all_vals.max()
+        margin = (vmax - vmin) * 0.05 if vmax > vmin else 0.5
+        ax.plot(
+            [vmin - margin, vmax + margin],
+            [vmin - margin, vmax + margin],
+            "k--",
+            alpha=0.4,
+            linewidth=1,
+            zorder=0,
+        )
+
+        ax.set_xlabel("True value")
+        ax.set_ylabel(f"Posterior {point_estimate}")
+        ax.set_title(
+            f"{label} — coverage: {empirical_cov:.1%} "
+            f"(expected: {ci_prob:.1%} ± {2 * binomial_se:.1%})",
+            fontsize=9,
+            wrap=True,
+        )
+        if mask_miss.any():
+            ax.legend(fontsize="small")
+
+    if created_figure:
+        fig.tight_layout()
+
+    if if_show:
+        plt.show()
+
+    return fig
+
+
+def plot_ecdf(sbc, if_show=True):
+    if not if_show:
+        matplotlib.use("Agg")
+    else:
+        matplotlib.use("TkAgg")
+
+    if sbc.method == "posterior":
+        fig = plot_ecdf_pit(
+            sbc.simulations,
+            group="posterior_sbc",
+            visuals={"xlabel": False},
+        )
+    else:
+        fig = plot_ecdf_pit(
+            sbc.simulations,
+            visuals={"xlabel": False},
+        )
+
+    if if_show:
+        plt.show()
+
+    return fig
diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
new file mode 100644
index 0000000..6bcdc00
--- /dev/null
+++ b/simuk/tests/test_plot.py
@@ -0,0 +1,232 @@
+import arviz_plots as azp
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import pymc as pm
+import pytest
+
+import simuk
+
+matplotlib.use("Agg")
+
+# Test data (same as test_prior_sbc.py)
+data = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
+sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
+
+with pm.Model() as centered_eight:
+    mu = pm.Normal("mu", mu=0, sigma=5)
+    tau = pm.HalfCauchy("tau", beta=5)
+    theta = pm.Normal("theta", mu=mu, sigma=tau, shape=8)
+    y_obs = pm.Normal("y", mu=theta, sigma=sigma, observed=data)
+
+
+@pytest.fixture(scope="module")
+def sbc_with_fits():
+    sbc = simuk.SBC(
+        centered_eight,
+        num_simulations=10,
+        sample_kwargs={"draws": 5, "tune": 5},
+    )
+    sbc.run_simulations()
+    return sbc
+
+
+@pytest.fixture(scope="module")
+def sbc_no_fits():
+    sbc = simuk.SBC(
+        centered_eight,
+        num_simulations=10,
+        sample_kwargs={"draws": 5, "tune": 5},
+        keep_fits=False,
+    )
+    sbc.run_simulations()
+    return sbc
+
+
+# Test data (same as test_posterior_sbc.py)
+default_rng = np.random.default_rng(1234)
+obs_data = default_rng.normal(2.0, 1.0, size=20)
+x_obs = np.linspace(0, 1, 20)
+y_obs_reg = 1.5 * x_obs + default_rng.normal(0, 0.5, size=20)
+
+
+with pm.Model() as simple_posterior_model:
+    mu = pm.Normal("mu", mu=0, sigma=5)
+    sigma = pm.HalfNormal("sigma", sigma=2)
+    y_data = pm.Data("y_data", obs_data)
+    pm.Normal("y", mu=mu, sigma=sigma, observed=y_data)
+
+with simple_posterior_model:
+    trace_simple = pm.sample(
+        draws=30,
+        tune=30,
+        chains=1,
+        random_seed=123,
+        progressbar=False,
+        compute_convergence_checks=False,
+    )
+
+
+@pytest.fixture(scope="module")
+def sbc_posterior_with_fits():
+    sbc = simuk.SBC(
+        simple_posterior_model,
+        trace=trace_simple,
+        method="posterior",
+        num_simulations=10,
+        sample_kwargs={"draws": 5, "tune": 5},
+    )
+    sbc.run_simulations()
+    return sbc
+
+
+@pytest.fixture(scope="module")
+def sbc_posterior_no_fits():
+    sbc = simuk.SBC(
+        simple_posterior_model,
+        trace=trace_simple,
+        method="posterior",
+        num_simulations=10,
+        sample_kwargs={"draws": 5, "tune": 5},
+        keep_fits=False,
+    )
+    sbc.run_simulations()
+    return sbc
+
+
+def test_ppr_requires_keep_fits(sbc_no_fits):
+    with pytest.raises(ValueError, match="keep_fits"):
+        fig = simuk.plot_parameter_recovery(sbc_no_fits, if_show=False)
+
+
+def test_ppr_requires_keep_fits_posterior(sbc_posterior_no_fits):
+    with pytest.raises(ValueError, match="keep_fits"):
+        fig = simuk.plot_parameter_recovery(sbc_posterior_no_fits, if_show=False)
+
+
+def test_ppr_requires_completed_simulations():
+    sbc = simuk.SBC(
+        centered_eight,
+        num_simulations=2,
+        sample_kwargs={"draws": 5, "tune": 5},
+    )
+    with pytest.raises(ValueError, match="No posteriors"):
+        fig = simuk.plot_parameter_recovery(sbc, if_show=False)
+
+
+def test_ppr_requires_completed_simulations_posterior():
+    sbc = simuk.SBC(
+        simple_posterior_model,
+        trace=trace_simple,
+        method="posterior",
+        num_simulations=2,
+        sample_kwargs={"draws": 5, "tune": 5},
+    )
+    with pytest.raises(ValueError, match="No posteriors"):
+        fig = simuk.plot_parameter_recovery(sbc, if_show=False)
+
+
+
+def test_ppr_basic(sbc_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
+def test_ppr_basic_posterior(sbc_posterior_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
+def test_ppr_var_names_filter(sbc_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, var_names=["mu"], if_show=False)
+    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
+    assert len(visible_axes) == 1
+
+
+def test_ppr_var_names_filter_posterior(sbc_posterior_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, var_names=["mu"], if_show=False)
+    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
+    assert len(visible_axes) == 1
+
+
+def test_ppr_with_transform(sbc_with_fits):
+    fig = simuk.plot_parameter_recovery(
+        sbc_with_fits,
+        transform=lambda name, val: np.mean(val),
+        if_show=False
+    )
+    assert isinstance(fig, plt.Figure)
+    # The mean transform reduces theta (8,) to scalar → 3 subplots total
+    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
+    assert len(visible_axes) == 3  # mu, tau, theta (all scalar after transform)
+
+
+def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
+    fig = simuk.plot_parameter_recovery(
+        sbc_posterior_with_fits,
+        transform=lambda name, val: np.mean(val),
+        if_show=False
+    )
+    assert isinstance(fig, plt.Figure)
+    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
+    assert len(visible_axes) == 2 # mu, sigma (all scalar after transform)
+
+
+def test_ppr_custom_ci_prob(sbc_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, ci_prob=0.5, if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
+def test_ppr_custom_ci_prob_posterior(sbc_posterior_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5, if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
+def test_ppr_with_preexisting_axes(sbc_with_fits):
+    # mu(1) + tau(1) + theta(8) = 10 subplots needed
+    fig, axes = plt.subplots(2, 5)
+    returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, axes=axes,if_show=False)
+    assert returned_fig is fig
+
+
+def test_ppr_with_preexisting_axes_posterior(sbc_posterior_with_fits):
+    # mu(1) + sigma(1) + theta(8) = 10 subplots needed
+    fig, axes = plt.subplots(2, 5)
+    returned_fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, axes=axes, if_show=False)
+    assert returned_fig is fig
+
+
+def test_ppr_median_point_estimate(sbc_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="median", if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
+def test_ppr_median_point_estimate_posterior(sbc_posterior_with_fits):
+    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="median", if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
+def test_ppr_invalid_point_estimate(sbc_with_fits):
+    with pytest.raises(ValueError, match="point_estimate"):
+        fig = simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="mode", if_show=False)
+
+
+def test_ppr_invalid_point_estimate_posterior(sbc_posterior_with_fits):
+    with pytest.raises(ValueError, match="point_estimate"):
+        fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
+
+
+def test_plot_ecdf_basic(sbc_with_fits, sbc_no_fits):
+    fig = simuk.plot_ecdf(sbc_with_fits, if_show=False)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
+
+    fig = simuk.plot_ecdf(sbc_no_fits, if_show=False)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
+
+
+def test_plot_ecdf_posterior(sbc_posterior_with_fits, sbc_posterior_no_fits):
+    fig = simuk.plot_ecdf(sbc_posterior_with_fits, if_show=False)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
+
+    fig = simuk.plot_ecdf(sbc_posterior_no_fits, if_show=False)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)

From 88d8dd5dec3999382fc257ed17db2c0ff09b7766 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Mon, 20 Jul 2026 17:25:37 +0200
Subject: [PATCH 04/16] refactor: move backend adapters to adapters folder

---
 simuk/adapters/__init__.py              | 0
 simuk/{ => adapters}/backend_adapter.py | 0
 simuk/{ => adapters}/numpyro_adapter.py | 2 +-
 simuk/{ => adapters}/pymc_adapter.py    | 2 +-
 simuk/sbc.py                            | 6 +++---
 5 files changed, 5 insertions(+), 5 deletions(-)
 create mode 100644 simuk/adapters/__init__.py
 rename simuk/{ => adapters}/backend_adapter.py (100%)
 rename simuk/{ => adapters}/numpyro_adapter.py (99%)
 rename simuk/{ => adapters}/pymc_adapter.py (99%)

diff --git a/simuk/adapters/__init__.py b/simuk/adapters/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/simuk/backend_adapter.py b/simuk/adapters/backend_adapter.py
similarity index 100%
rename from simuk/backend_adapter.py
rename to simuk/adapters/backend_adapter.py
diff --git a/simuk/numpyro_adapter.py b/simuk/adapters/numpyro_adapter.py
similarity index 99%
rename from simuk/numpyro_adapter.py
rename to simuk/adapters/numpyro_adapter.py
index 2cc4eee..fe45587 100644
--- a/simuk/numpyro_adapter.py
+++ b/simuk/adapters/numpyro_adapter.py
@@ -9,7 +9,7 @@
 from numpyro.handlers import seed, trace
 from numpyro.infer import MCMC, Predictive
 
-from simuk.backend_adapter import BackendAdapter
+from simuk.adapters.backend_adapter import BackendAdapter
 
 log = logging.getLogger(__name__)
 
diff --git a/simuk/pymc_adapter.py b/simuk/adapters/pymc_adapter.py
similarity index 99%
rename from simuk/pymc_adapter.py
rename to simuk/adapters/pymc_adapter.py
index 91dd553..d3bb2c0 100644
--- a/simuk/pymc_adapter.py
+++ b/simuk/adapters/pymc_adapter.py
@@ -7,7 +7,7 @@
 import xarray as xr
 from arviz_base import dict_to_dataset, extract
 
-from simuk.backend_adapter import BackendAdapter
+from simuk.adapters.backend_adapter import BackendAdapter
 
 
 class PymcAdapter(BackendAdapter):
diff --git a/simuk/sbc.py b/simuk/sbc.py
index da7e89d..575c345 100644
--- a/simuk/sbc.py
+++ b/simuk/sbc.py
@@ -219,13 +219,13 @@ def __init__(
         self._seeds = self._get_seeds()
 
         if hasattr(model, "basic_RVs") and isinstance(model, pm.Model):
-            from simuk.pymc_adapter import PymcAdapter  # noqa: PLC0415
+            from simuk.adapters.pymc_adapter import PymcAdapter  # noqa: PLC0415
 
             self.engine = "pymc"
             self.model = model
             self.adapter = PymcAdapter(self.model, simulator, trace, augment_observed, update_data)
         elif hasattr(model, "formula"):
-            from simuk.pymc_adapter import PymcAdapter  # noqa: PLC0415
+            from simuk.adapters.pymc_adapter import PymcAdapter  # noqa: PLC0415
 
             self.engine = "bambi"
             model.build()
@@ -236,7 +236,7 @@ def __init__(
             self.adapter = PymcAdapter(self.model, simulator, trace, augment_observed, update_data)
         elif isinstance(model, MCMCKernel):
             # runtime import so an environment with only Pymc can run SBC over Pymc models.
-            from simuk.numpyro_adapter import NumpyroAdapter  # noqa: PLC0415
+            from simuk.adapters.numpyro_adapter import NumpyroAdapter  # noqa: PLC0415
 
             self.engine = "numpyro"
             self.numpyro_model = model

From 76b0d1fece9a091823af844508e71739654af2fa Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Mon, 20 Jul 2026 18:03:54 +0200
Subject: [PATCH 05/16] fix: linting errors

---
 simuk/plots.py           |  3 ++-
 simuk/tests/test_plot.py | 29 +++++++++++++----------------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/simuk/plots.py b/simuk/plots.py
index f06b43e..22d1ecd 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -3,8 +3,9 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
-from arviz_plots import plot_ecdf_pit
+from arviz_plots import plot_ecdf_pit, style
 
+style.use("arviz-variat")
 
 def _extract_draws(posterior, name):
     """Extract posterior draws for a parameter as a numpy array.
diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index 6bcdc00..8ba029c 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -96,12 +96,12 @@ def sbc_posterior_no_fits():
 
 def test_ppr_requires_keep_fits(sbc_no_fits):
     with pytest.raises(ValueError, match="keep_fits"):
-        fig = simuk.plot_parameter_recovery(sbc_no_fits, if_show=False)
+        simuk.plot_parameter_recovery(sbc_no_fits, if_show=False)
 
 
 def test_ppr_requires_keep_fits_posterior(sbc_posterior_no_fits):
     with pytest.raises(ValueError, match="keep_fits"):
-        fig = simuk.plot_parameter_recovery(sbc_posterior_no_fits, if_show=False)
+        simuk.plot_parameter_recovery(sbc_posterior_no_fits, if_show=False)
 
 
 def test_ppr_requires_completed_simulations():
@@ -111,7 +111,7 @@ def test_ppr_requires_completed_simulations():
         sample_kwargs={"draws": 5, "tune": 5},
     )
     with pytest.raises(ValueError, match="No posteriors"):
-        fig = simuk.plot_parameter_recovery(sbc, if_show=False)
+        simuk.plot_parameter_recovery(sbc, if_show=False)
 
 
 def test_ppr_requires_completed_simulations_posterior():
@@ -123,8 +123,7 @@ def test_ppr_requires_completed_simulations_posterior():
         sample_kwargs={"draws": 5, "tune": 5},
     )
     with pytest.raises(ValueError, match="No posteriors"):
-        fig = simuk.plot_parameter_recovery(sbc, if_show=False)
-
+        simuk.plot_parameter_recovery(sbc, if_show=False)
 
 
 def test_ppr_basic(sbc_with_fits):
@@ -151,9 +150,7 @@ def test_ppr_var_names_filter_posterior(sbc_posterior_with_fits):
 
 def test_ppr_with_transform(sbc_with_fits):
     fig = simuk.plot_parameter_recovery(
-        sbc_with_fits,
-        transform=lambda name, val: np.mean(val),
-        if_show=False
+        sbc_with_fits, transform=lambda name, val: np.mean(val), if_show=False
     )
     assert isinstance(fig, plt.Figure)
     # The mean transform reduces theta (8,) to scalar → 3 subplots total
@@ -163,13 +160,11 @@ def test_ppr_with_transform(sbc_with_fits):
 
 def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(
-        sbc_posterior_with_fits,
-        transform=lambda name, val: np.mean(val),
-        if_show=False
+        sbc_posterior_with_fits, transform=lambda name, val: np.mean(val), if_show=False
     )
     assert isinstance(fig, plt.Figure)
     visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
-    assert len(visible_axes) == 2 # mu, sigma (all scalar after transform)
+    assert len(visible_axes) == 2  # mu, sigma (all scalar after transform)
 
 
 def test_ppr_custom_ci_prob(sbc_with_fits):
@@ -185,7 +180,7 @@ def test_ppr_custom_ci_prob_posterior(sbc_posterior_with_fits):
 def test_ppr_with_preexisting_axes(sbc_with_fits):
     # mu(1) + tau(1) + theta(8) = 10 subplots needed
     fig, axes = plt.subplots(2, 5)
-    returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, axes=axes,if_show=False)
+    returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, axes=axes, if_show=False)
     assert returned_fig is fig
 
 
@@ -202,18 +197,20 @@ def test_ppr_median_point_estimate(sbc_with_fits):
 
 
 def test_ppr_median_point_estimate_posterior(sbc_posterior_with_fits):
-    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="median", if_show=False)
+    fig = simuk.plot_parameter_recovery(
+        sbc_posterior_with_fits, point_estimate="median", if_show=False
+    )
     assert isinstance(fig, plt.Figure)
 
 
 def test_ppr_invalid_point_estimate(sbc_with_fits):
     with pytest.raises(ValueError, match="point_estimate"):
-        fig = simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="mode", if_show=False)
+        simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="mode", if_show=False)
 
 
 def test_ppr_invalid_point_estimate_posterior(sbc_posterior_with_fits):
     with pytest.raises(ValueError, match="point_estimate"):
-        fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
+        simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
 
 
 def test_plot_ecdf_basic(sbc_with_fits, sbc_no_fits):

From b758f99b46dae3d07fa0a79b0dc9c7912823e907 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Mon, 20 Jul 2026 18:04:59 +0200
Subject: [PATCH 06/16] docs: fix docs to match new plots

---
 docs/examples/gallery/posterior_sbc.ipynb | 22 ++++++++------------
 docs/examples/gallery/prior_sbc.ipynb     | 25 ++++++++---------------
 docs/index.rst                            |  7 ++-----
 3 files changed, 18 insertions(+), 36 deletions(-)

diff --git a/docs/examples/gallery/posterior_sbc.ipynb b/docs/examples/gallery/posterior_sbc.ipynb
index 1caa14b..a329a37 100644
--- a/docs/examples/gallery/posterior_sbc.ipynb
+++ b/docs/examples/gallery/posterior_sbc.ipynb
@@ -24,19 +24,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "import pymc as pm\n",
-    "from arviz_plots import plot_ecdf_pit, style\n",
     "\n",
     "import simuk\n",
     "\n",
     "random_seed = 42\n",
-    "rng = np.random.default_rng(random_seed)\n",
-    "style.use(\"arviz-variat\")"
+    "rng = np.random.default_rng(random_seed)"
    ]
   },
   {
@@ -228,7 +226,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -243,11 +241,7 @@
     }
    ],
    "source": [
-    "plot_ecdf_pit(\n",
-    "    sbc.simulations,\n",
-    "    group=\"posterior_sbc\",\n",
-    "    visuals={\"xlabel\": False},\n",
-    ");"
+    "simuk.plot_ecdf(sbc)"
    ]
   },
   {
@@ -311,7 +305,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -336,7 +330,7 @@
     }
    ],
    "source": [
-    "plot_ecdf_pit(skewed_sbc.simulations, group=\"posterior_sbc\", visuals={\"xlabel\": False})"
+    "simuk.plot_ecdf(skewed_sbc)"
    ]
   },
   {
@@ -348,7 +342,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -364,7 +358,7 @@
    ],
    "source": [
     "sbc.compute_rank_statistics(lambda _, param_value: param_value)\n",
-    "plot_ecdf_pit(sbc.simulations, group=\"posterior_sbc\", visuals={\"xlabel\": False})"
+    "simuk.plot_ecdf(sbc)"
    ]
   },
   {
diff --git a/docs/examples/gallery/prior_sbc.ipynb b/docs/examples/gallery/prior_sbc.ipynb
index 5ece4ee..4466c88 100644
--- a/docs/examples/gallery/prior_sbc.ipynb
+++ b/docs/examples/gallery/prior_sbc.ipynb
@@ -9,16 +9,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "from arviz_plots import plot_ecdf_pit, style\n",
-    "\n",
-    "import simuk\n",
     "\n",
-    "style.use(\"arviz-variat\")"
+    "import simuk"
    ]
   },
   {
@@ -89,7 +86,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -104,10 +101,7 @@
     }
    ],
    "source": [
-    "plot_ecdf_pit(\n",
-    "    sbc.simulations,\n",
-    "    visuals={\"xlabel\": False},\n",
-    ");"
+    "simuk.plot_ecdf(sbc)"
    ]
   },
   {
@@ -163,7 +157,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -188,7 +182,7 @@
     }
    ],
    "source": [
-    "plot_ecdf_pit(sbc.simulations)"
+    "simuk.plot_ecdf(sbc)"
    ]
   },
   {
@@ -269,7 +263,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -284,10 +278,7 @@
     }
    ],
    "source": [
-    "plot_ecdf_pit(\n",
-    "    sbc.simulations,\n",
-    "    visuals={\"xlabel\": False},\n",
-    ");"
+    "simuk.plot_ecdf(sbc)"
    ]
   },
   {
diff --git a/docs/index.rst b/docs/index.rst
index e5b8bc0..aa2979d 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -23,7 +23,6 @@ In our case, we will take a PyMC model and pass it into our ``SBC`` class.
 
 .. code-block:: python
 
-    from arviz_plots import plot_ecdf_pit
     import numpy as np
     import pymc as pm
 
@@ -50,9 +49,7 @@ Plot the empirical CDF to compare the differences between the prior and posterio
 
 .. code-block:: python
 
-    plot_ecdf_pit(sbc.simulations,
-                visuals={"xlabel":False},
-    );
+    simuk.plot_ecdf(sbc)
 
 The lines should be nearly uniform and fall within the oval envelope. It suggests that the prior and posterior distributions
 are properly aligned and that there are no significant biases or issues with the model.
@@ -118,7 +115,7 @@ Currently, it's only implemented for PyMC.
     )
     post_sbc.run_simulations()
 
-    plot_ecdf_pit(post_sbc.simulations, group="posterior_sbc", visuals={"xlabel": False})
+    simuk.plot_ecdf(post_sbc)
 
 For more advanced use cases, such as custom data augmentation or re-evaluating rank statistics, check out the :doc:`Posterior SBC tutorial <examples/gallery/posterior_sbc>`.
 

From dfa1ec6a51050c30c452fb6262b7d1f8a7580cd0 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Mon, 20 Jul 2026 22:37:07 +0200
Subject: [PATCH 07/16] fix: linting errors

---
 simuk/plots.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/simuk/plots.py b/simuk/plots.py
index 22d1ecd..a30d19a 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -7,6 +7,7 @@
 
 style.use("arviz-variat")
 
+
 def _extract_draws(posterior, name):
     """Extract posterior draws for a parameter as a numpy array.
 

From d8171d8120f86ebc5eadfd050dea6c62ed5b5538 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Tue, 21 Jul 2026 12:19:06 +0200
Subject: [PATCH 08/16] tests: add seed

---
 simuk/tests/test_plot.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index 8ba029c..a1189da 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -25,7 +25,8 @@ def sbc_with_fits():
     sbc = simuk.SBC(
         centered_eight,
         num_simulations=10,
-        sample_kwargs={"draws": 5, "tune": 5},
+        sample_kwargs={"draws": 10, "tune": 10},
+        seed=42,
     )
     sbc.run_simulations()
     return sbc
@@ -36,8 +37,9 @@ def sbc_no_fits():
     sbc = simuk.SBC(
         centered_eight,
         num_simulations=10,
-        sample_kwargs={"draws": 5, "tune": 5},
+        sample_kwargs={"draws": 10, "tune": 10},
         keep_fits=False,
+        seed=42,
     )
     sbc.run_simulations()
     return sbc
@@ -74,6 +76,7 @@ def sbc_posterior_with_fits():
         trace=trace_simple,
         method="posterior",
         num_simulations=10,
+        seed=42,
         sample_kwargs={"draws": 5, "tune": 5},
     )
     sbc.run_simulations()
@@ -88,6 +91,7 @@ def sbc_posterior_no_fits():
         method="posterior",
         num_simulations=10,
         sample_kwargs={"draws": 5, "tune": 5},
+        seed=42,
         keep_fits=False,
     )
     sbc.run_simulations()

From d232b32376b3a3cf2a281853b5b6331da1daea25 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Tue, 21 Jul 2026 22:48:53 +0200
Subject: [PATCH 09/16] tests: unkown pytest problem

---
 simuk/tests/conftest.py | 10 ++++++++++
 1 file changed, 10 insertions(+)
 create mode 100644 simuk/tests/conftest.py

diff --git a/simuk/tests/conftest.py b/simuk/tests/conftest.py
new file mode 100644
index 0000000..4eee837
--- /dev/null
+++ b/simuk/tests/conftest.py
@@ -0,0 +1,10 @@
+# This is needed as running all tests together now causes the test to fail
+# when the same doesn't happen when running tests separately.
+# The Error is the same: AttributeError: module 'numpyro.infer' has no
+# attribute 'initialization' in arviz_base.io_numpyro
+
+# This may have something to do with the lazy loading used in arviz_base.io_numpyro,
+# in accordance with import ordering and different behaviors when using from-imports.
+# However, I'm unable to recreate it using a minimal example. This at least enforces
+# the numpyro.infer to have the initialization attribute.
+import numpyro.infer.initialization  # noqa: F401

From 8dd653ba735e98d10e3cab71481c1a4763e82331 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:27:43 +0200
Subject: [PATCH 10/16] tests: more tests

---
 simuk/plots.py           |  5 +----
 simuk/tests/test_plot.py | 24 ++++++++++++++++++++----
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/simuk/plots.py b/simuk/plots.py
index a30d19a..a5da1f4 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -127,10 +127,7 @@ def plot_parameter_recovery(
     posterior draws per fit is small, the actual uncertainty in empirical
     coverage will be larger than the displayed ± range.
     """
-    if not if_show:
-        matplotlib.use("Agg")
-    else:
-        matplotlib.use("TkAgg")
+
 
     # ---- validation --------------------------------------------------------
     if not sbc.keep_fits:
diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index a1189da..e442c9d 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -161,6 +161,12 @@ def test_ppr_with_transform(sbc_with_fits):
     visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
     assert len(visible_axes) == 3  # mu, tau, theta (all scalar after transform)
 
+def test_ppr_with_bad_transform(sbc_with_fits):
+    with pytest.raises(ValueError, match="`transform` should be a function or None"):
+        simuk.plot_parameter_recovery(
+            sbc_with_fits, transform="bad transform", if_show=False
+        )
+
 
 def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(
@@ -172,12 +178,15 @@ def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
 
 
 def test_ppr_custom_ci_prob(sbc_with_fits):
-    fig = simuk.plot_parameter_recovery(sbc_with_fits, ci_prob=0.5, if_show=False)
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, ci_prob=0.5)
+    plt.close(fig)
     assert isinstance(fig, plt.Figure)
 
 
+
 def test_ppr_custom_ci_prob_posterior(sbc_posterior_with_fits):
-    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5, if_show=False)
+    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5)
+    plt.close(fig)
     assert isinstance(fig, plt.Figure)
 
 
@@ -187,6 +196,12 @@ def test_ppr_with_preexisting_axes(sbc_with_fits):
     returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, axes=axes, if_show=False)
     assert returned_fig is fig
 
+def test_ppr_with_insufficient_preexisting_axes(sbc_with_fits):
+    # mu(1) + tau(1) + theta(8) = 10 subplots needed
+    with pytest.raises(ValueError, match="axes but only"):
+        _, axes = plt.subplots(2, 4)
+        simuk.plot_parameter_recovery(sbc_with_fits, axes=axes, if_show=False)
+
 
 def test_ppr_with_preexisting_axes_posterior(sbc_posterior_with_fits):
     # mu(1) + sigma(1) + theta(8) = 10 subplots needed
@@ -216,7 +231,6 @@ def test_ppr_invalid_point_estimate_posterior(sbc_posterior_with_fits):
     with pytest.raises(ValueError, match="point_estimate"):
         simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
 
-
 def test_plot_ecdf_basic(sbc_with_fits, sbc_no_fits):
     fig = simuk.plot_ecdf(sbc_with_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
@@ -229,5 +243,7 @@ def test_plot_ecdf_posterior(sbc_posterior_with_fits, sbc_posterior_no_fits):
     fig = simuk.plot_ecdf(sbc_posterior_with_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
 
-    fig = simuk.plot_ecdf(sbc_posterior_no_fits, if_show=False)
+    plt.ion()
+    fig = simuk.plot_ecdf(sbc_posterior_no_fits)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
+    plt.close("all")

From e8a7d1d6d1f80a9409462a0fa30541ffe4793dd8 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:30:25 +0200
Subject: [PATCH 11/16] chore: linting

---
 simuk/plots.py           | 2 --
 simuk/tests/test_plot.py | 8 ++++----
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/simuk/plots.py b/simuk/plots.py
index a5da1f4..862e441 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -127,8 +127,6 @@ def plot_parameter_recovery(
     posterior draws per fit is small, the actual uncertainty in empirical
     coverage will be larger than the displayed ± range.
     """
-
-
     # ---- validation --------------------------------------------------------
     if not sbc.keep_fits:
         raise ValueError(
diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index e442c9d..98284ab 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -161,11 +161,10 @@ def test_ppr_with_transform(sbc_with_fits):
     visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
     assert len(visible_axes) == 3  # mu, tau, theta (all scalar after transform)
 
+
 def test_ppr_with_bad_transform(sbc_with_fits):
     with pytest.raises(ValueError, match="`transform` should be a function or None"):
-        simuk.plot_parameter_recovery(
-            sbc_with_fits, transform="bad transform", if_show=False
-        )
+        simuk.plot_parameter_recovery(sbc_with_fits, transform="bad transform", if_show=False)
 
 
 def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
@@ -183,7 +182,6 @@ def test_ppr_custom_ci_prob(sbc_with_fits):
     assert isinstance(fig, plt.Figure)
 
 
-
 def test_ppr_custom_ci_prob_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5)
     plt.close(fig)
@@ -196,6 +194,7 @@ def test_ppr_with_preexisting_axes(sbc_with_fits):
     returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, axes=axes, if_show=False)
     assert returned_fig is fig
 
+
 def test_ppr_with_insufficient_preexisting_axes(sbc_with_fits):
     # mu(1) + tau(1) + theta(8) = 10 subplots needed
     with pytest.raises(ValueError, match="axes but only"):
@@ -231,6 +230,7 @@ def test_ppr_invalid_point_estimate_posterior(sbc_posterior_with_fits):
     with pytest.raises(ValueError, match="point_estimate"):
         simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
 
+
 def test_plot_ecdf_basic(sbc_with_fits, sbc_no_fits):
     fig = simuk.plot_ecdf(sbc_with_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)

From b4eefd7feb7c43cf5d4b253981d56355534a6d71 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:38:48 +0200
Subject: [PATCH 12/16] tests: fix test_plot_ecdf_posterior

---
 simuk/tests/test_plot.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index 98284ab..1b354f4 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -243,7 +243,5 @@ def test_plot_ecdf_posterior(sbc_posterior_with_fits, sbc_posterior_no_fits):
     fig = simuk.plot_ecdf(sbc_posterior_with_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
 
-    plt.ion()
-    fig = simuk.plot_ecdf(sbc_posterior_no_fits)
+    fig = simuk.plot_ecdf(sbc_posterior_no_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
-    plt.close("all")

From d5fb96744a52b1a3eb271b7196a78b7b39ef3dc1 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Wed, 22 Jul 2026 00:23:07 +0200
Subject: [PATCH 13/16] tests: more tests

---
 simuk/plots.py           |  5 ----
 simuk/tests/test_plot.py | 61 ++++++++++++++++++++++++++++++++--------
 2 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/simuk/plots.py b/simuk/plots.py
index 862e441..c8f2ccc 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -1,6 +1,5 @@
 """Parameter recovery plotting for simulation-based calibration."""
 
-import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 from arviz_plots import plot_ecdf_pit, style
@@ -299,10 +298,6 @@ def plot_parameter_recovery(
 
 
 def plot_ecdf(sbc, if_show=True):
-    if not if_show:
-        matplotlib.use("Agg")
-    else:
-        matplotlib.use("TkAgg")
 
     if sbc.method == "posterior":
         fig = plot_ecdf_pit(
diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index 1b354f4..bb0647c 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -1,23 +1,31 @@
 import arviz_plots as azp
-import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+import numpyro
+import numpyro.distributions as dist
 import pymc as pm
 import pytest
+from numpyro.infer import NUTS
 
 import simuk
 
-matplotlib.use("Agg")
-
 # Test data (same as test_prior_sbc.py)
-data = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
-sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
+plot_data = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
+sigma_eight_schools = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
 
 with pm.Model() as centered_eight:
     mu = pm.Normal("mu", mu=0, sigma=5)
     tau = pm.HalfCauchy("tau", beta=5)
     theta = pm.Normal("theta", mu=mu, sigma=tau, shape=8)
-    y_obs = pm.Normal("y", mu=theta, sigma=sigma, observed=data)
+    y_obs = pm.Normal("y", mu=theta, sigma=sigma_eight_schools, observed=plot_data)
+
+
+def eight_schools_cauchy_prior(J, sigma, y=None):
+    mu = numpyro.sample("mu", dist.Normal(0, 5))
+    tau = numpyro.sample("tau", dist.HalfCauchy(5))
+    with numpyro.plate("J", J):
+        theta = numpyro.sample("theta", dist.Normal(mu, tau))
+    numpyro.sample("y", dist.Normal(theta, sigma), obs=y)
 
 
 @pytest.fixture(scope="module")
@@ -32,6 +40,19 @@ def sbc_with_fits():
     return sbc
 
 
+@pytest.fixture(scope="module")
+def sbc_with_fits_numpyro():
+    sbc = simuk.SBC(
+        NUTS(eight_schools_cauchy_prior),
+        data_dir={"J": 8, "sigma": sigma_eight_schools, "y": plot_data},
+        num_simulations=10,
+        sample_kwargs={"num_warmup": 10, "num_samples": 10},
+        seed=42,
+    )
+    sbc.run_simulations()
+    return sbc
+
+
 @pytest.fixture(scope="module")
 def sbc_no_fits():
     sbc = simuk.SBC(
@@ -54,9 +75,9 @@ def sbc_no_fits():
 
 with pm.Model() as simple_posterior_model:
     mu = pm.Normal("mu", mu=0, sigma=5)
-    sigma = pm.HalfNormal("sigma", sigma=2)
+    sigma_pymc = pm.HalfNormal("sigma", sigma=2)
     y_data = pm.Data("y_data", obs_data)
-    pm.Normal("y", mu=mu, sigma=sigma, observed=y_data)
+    pm.Normal("y", mu=mu, sigma=sigma_pymc, observed=y_data)
 
 with simple_posterior_model:
     trace_simple = pm.sample(
@@ -135,6 +156,11 @@ def test_ppr_basic(sbc_with_fits):
     assert isinstance(fig, plt.Figure)
 
 
+def test_plot_ppr_basic_numpyro(sbc_with_fits_numpyro):
+    fig = simuk.plot_parameter_recovery(sbc_with_fits_numpyro, if_show=False)
+    assert isinstance(fig, plt.Figure)
+
+
 def test_ppr_basic_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, if_show=False)
     assert isinstance(fig, plt.Figure)
@@ -177,14 +203,13 @@ def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
 
 
 def test_ppr_custom_ci_prob(sbc_with_fits):
-    fig = simuk.plot_parameter_recovery(sbc_with_fits, ci_prob=0.5)
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, ci_prob=0.5, if_show=False)
     plt.close(fig)
     assert isinstance(fig, plt.Figure)
 
 
 def test_ppr_custom_ci_prob_posterior(sbc_posterior_with_fits):
-    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5)
-    plt.close(fig)
+    fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5, if_show=False)
     assert isinstance(fig, plt.Figure)
 
 
@@ -231,6 +256,20 @@ def test_ppr_invalid_point_estimate_posterior(sbc_posterior_with_fits):
         simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
 
 
+def test_ppr_show_branch(monkeypatch, sbc_with_fits):
+    called = {"value": False}
+
+    def fake_show():
+        called["value"] = True
+
+    monkeypatch.setattr(plt, "show", fake_show)
+
+    fig = simuk.plot_parameter_recovery(sbc_with_fits, if_show=True)
+    assert isinstance(fig, plt.Figure)
+    assert called["value"]
+    plt.close(fig)
+
+
 def test_plot_ecdf_basic(sbc_with_fits, sbc_no_fits):
     fig = simuk.plot_ecdf(sbc_with_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)

From e861078de5411fc3550463f1f8d3d57cce4671fe Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Wed, 22 Jul 2026 00:29:27 +0200
Subject: [PATCH 14/16] tests: more tests

---
 simuk/tests/test_plot.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index bb0647c..9101f94 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -284,3 +284,17 @@ def test_plot_ecdf_posterior(sbc_posterior_with_fits, sbc_posterior_no_fits):
 
     fig = simuk.plot_ecdf(sbc_posterior_no_fits, if_show=False)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
+
+
+def test_ecdf_show_branch(monkeypatch, sbc_with_fits):
+    called = {"value": False}
+
+    def fake_show():
+        called["value"] = True
+
+    monkeypatch.setattr(plt, "show", fake_show)
+
+    fig = simuk.plot_ecdf(sbc_with_fits, if_show=True)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
+    assert called["value"]
+    plt.close("all")

From ce44c7c99b009d7b79439dcc1afe27bae9835b9d Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Sat, 25 Jul 2026 19:10:49 +0800
Subject: [PATCH 15/16] Update simuk/plots.py

Co-authored-by: Osvaldo A Martin <aloctavodia@gmail.com>
---
 simuk/plots.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/simuk/plots.py b/simuk/plots.py
index c8f2ccc..e8f3825 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -84,8 +84,7 @@ def plot_parameter_recovery(
         ``keep_fits=True``.
     ci_prob : float, default 0.94
         Probability mass of the credible interval (e.g. 0.94 for a 94 %
-        CI).  The interval is computed from posterior quantiles — no
-        distributional assumptions (e.g. normality) are made.
+        CI).  The interval is computed from posterior quantiles.
     point_estimate : {"mean", "median"}, default "mean"
         Which posterior point estimate to plot on the y-axis.
     transform : callable, optional

From 7403c57baf1aa10539e278e45ee90b40a9cb7580 Mon Sep 17 00:00:00 2001
From: cab14bacc <86755693+Cab14bacc@users.noreply.github.com>
Date: Tue, 28 Jul 2026 00:03:29 +0200
Subject: [PATCH 16/16] feat: make plot more arvizian

---
 simuk/plots.py           | 547 +++++++++++++++++++++------------------
 simuk/tests/conftest.py  |  10 -
 simuk/tests/test_plot.py | 104 +++-----
 3 files changed, 331 insertions(+), 330 deletions(-)
 delete mode 100644 simuk/tests/conftest.py

diff --git a/simuk/plots.py b/simuk/plots.py
index e8f3825..6480756 100644
--- a/simuk/plots.py
+++ b/simuk/plots.py
@@ -1,163 +1,70 @@
 """Parameter recovery plotting for simulation-based calibration."""
 
-import matplotlib.pyplot as plt
+from importlib import import_module
+
+import arviz_plots as azp
 import numpy as np
+import xarray as xr
+from arviz_base import rcParams
+from arviz_base.validate import validate_dict_argument, validate_or_use_rcparam
 from arviz_plots import plot_ecdf_pit, style
+from arviz_plots.plot_collection import PlotCollection, backend_from_object
+from arviz_plots.plots.utils import get_visual_kwargs
 
 style.use("arviz-variat")
 
 
-def _extract_draws(posterior, name):
-    """Extract posterior draws for a parameter as a numpy array.
-
-    Handles both PyMC-style posteriors (with a ``sample`` dimension from
-    ``arviz_base.extract``) and NumPyro-style posteriors (with ``chain``
-    and ``draw`` dimensions from ``arviz_base.from_numpyro``).
-
-    Parameters
-    ----------
-    posterior : xarray.Dataset
-        Posterior draws for a single SBC iteration.
-    name : str
-        Parameter name.
-
-    Returns
-    -------
-    numpy.ndarray
-        Array with shape ``(n_draws, *param_shape)``.
-    """
-    var = posterior[name]
-    if "sample" in var.dims:
-        return np.asarray(var.values)
-    else:
-        stacked = var.stack(__sample__=("chain", "draw"))
-        return np.asarray(stacked.transpose("__sample__", ...).values)
-
-
-def _extract_true_value(ref_params, name, idx):
-    """Extract the true (reference) parameter value for a simulation.
-
-    Parameters
-    ----------
-    ref_params : xarray.Dataset or dict
-        Reference parameters. An xarray Dataset (PyMC) or a plain dict of
-        arrays (NumPyro).
-    name : str
-        Parameter name.
-    idx : int
-        Simulation index.
-
-    Returns
-    -------
-    numpy.ndarray
-        The true parameter value (0-d for scalars, 1-d+ for vectors).
-    """
-    if hasattr(ref_params, "data_vars"):
-        return np.asarray(ref_params[name].isel(sample=idx).values)
+def plot_ecdf(sbc, if_show=True):
+    """Plot the empirical CDF of the SBC simulations."""
+    if sbc.method == "posterior":
+        fig = plot_ecdf_pit(
+            sbc.simulations,
+            group="posterior_sbc",
+            visuals={"xlabel": False},
+        )
     else:
-        return np.asarray(ref_params[name][idx])
-
-
-def plot_parameter_recovery(
-    sbc,
-    ci_prob=0.94,
-    point_estimate="mean",
-    transform=None,
-    var_names=None,
-    figsize=None,
-    axes=None,
-    if_show=True,
-):
-    r"""Plot parameter recovery: posterior point estimate vs true simulated value.
-
-    For each SBC iteration the true (reference) parameter value is plotted
-    on the x-axis and the posterior point estimate (mean or median) on the
-    y-axis, with credible-interval error bars.  Points whose credible
-    interval covers the true value are coloured blue; misses are red.
-
-    A *y = x* diagonal reference line indicates perfect recovery.
-
-    Parameters
-    ----------
-    sbc : simuk.SBC
-        An SBC instance that has completed simulations with
-        ``keep_fits=True``.
-    ci_prob : float, default 0.94
-        Probability mass of the credible interval (e.g. 0.94 for a 94 %
-        CI).  The interval is computed from posterior quantiles.
-    point_estimate : {"mean", "median"}, default "mean"
-        Which posterior point estimate to plot on the y-axis.
-    transform : callable, optional
-        ``(param_name, param_value) -> transformed_value``.  Applied to
-        both true values and posterior draws before summarising.  Defaults
-        to the transform set at SBC initialisation.
-    var_names : list[str], optional
-        Subset of parameter names to plot.  Defaults to all free
-        parameters.
-    figsize : tuple[float, float], optional
-        Figure size in inches.
-    axes : matplotlib Axes or array of Axes, optional
-        Pre-existing axes to draw on (one per scalar subplot).
+        fig = plot_ecdf_pit(
+            sbc.simulations,
+            group="prior_sbc",
+            visuals={"xlabel": False},
+        )
 
-    Returns
-    -------
-    matplotlib.figure.Figure
-        The figure containing the parameter recovery subplot(s).
+    if if_show:
+        fig.show()
 
-    Notes
-    -----
-    **Credible intervals** are quantile-based: for a 94 % CI the lower
-    and upper bounds are the 3rd and 97th percentiles of the posterior
-    draws.
+    return fig
 
-    **Coverage** reported in subplot titles is the empirical fraction of
-    simulations whose CI contains the true value.  The "expected" range
-    uses the binomial standard error
 
-    .. math::
+def _build_recovery_dataset(sbc, ci_prob, point_estimate, transform, var_names):
+    """Build an xarray.Dataset suitable for arviz_plots faceting.
 
-        \mathrm{SE} = \sqrt{\frac{p\,(1-p)}{N}}
+    The dataset contains a single variable ``recovery`` with dimensions
+    ``(simulation, parameter, quantity)``.  The ``quantity`` coordinate has
+    the values ``true``, ``estimate``, ``ci_low`` and ``ci_high``.  A
+    companion boolean variable ``covered`` indicates whether each
+    simulation's true value falls inside the credible interval.
 
-    which accounts only for the sampling noise from having a finite number
-    *N* of SBC iterations.  This is a **lower bound** on the true
-    uncertainty: it ignores the additional noise introduced by estimating
-    each CI from a finite number of MCMC draws.  When the number of
-    posterior draws per fit is small, the actual uncertainty in empirical
-    coverage will be larger than the displayed ± range.
+    Using a flat ``parameter`` dimension avoids naming conflicts between
+    data variables and their dimension coordinates when parameters are
+    scalar.
     """
-    # ---- validation --------------------------------------------------------
-    if not sbc.keep_fits:
-        raise ValueError(
-            "`plot_parameter_recovery` requires the SBC instance to have "
-            "been run with `keep_fits=True`."
-        )
-    if not sbc.posteriors:
-        raise ValueError("No posteriors found. Call `sbc.run_simulations()` before plotting.")
-    if point_estimate not in ("mean", "median"):
-        raise ValueError(f"`point_estimate` must be 'mean' or 'median', got '{point_estimate}'")
-    if transform is None:
-        transform = sbc._transform
-    elif not callable(transform):
-        raise ValueError("`transform` should be a function or None")
-
     ref_params = sbc.kept_simulation_params.ref_params
-    if var_names is None:
-        var_names = list(sbc.kept_simulation_params.var_names)
-
     n_sims = len(sbc.posteriors)
     alpha_lo = (1 - ci_prob) / 2
     alpha_hi = 1 - alpha_lo
 
-    # ---- compute summaries -------------------------------------------------
-    summaries = {}
+    param_labels = []
+    summaries_list = []
+    covered_list = []
+
     for name in var_names:
         true_list, est_list, lo_list, hi_list = [], [], [], []
 
         for idx in range(n_sims):
             posterior = sbc.posteriors[idx]
-            draws = _extract_draws(posterior, name)
-            transformed_draws = np.array([transform(name, d) for d in draws])
+            draws = posterior[name].transpose("sample", ...).values
 
+            transformed_draws = np.array([transform(name, d) for d in draws])
             if point_estimate == "mean":
                 est = np.mean(transformed_draws, axis=0)
             else:
@@ -166,7 +73,7 @@ def plot_parameter_recovery(
             ci_lo = np.quantile(transformed_draws, alpha_lo, axis=0)
             ci_hi = np.quantile(transformed_draws, alpha_hi, axis=0)
 
-            true_val = _extract_true_value(ref_params, name, idx)
+            true_val = ref_params[name].isel(sample=idx).values
             transformed_true = np.asarray(transform(name, true_val), dtype=float)
 
             true_list.append(np.atleast_1d(transformed_true))
@@ -174,143 +81,271 @@ def plot_parameter_recovery(
             lo_list.append(np.atleast_1d(ci_lo))
             hi_list.append(np.atleast_1d(ci_hi))
 
-        summaries[name] = {
-            "true_values": np.array(true_list),
-            "point_estimates": np.array(est_list),
-            "ci_lower": np.array(lo_list),
-            "ci_upper": np.array(hi_list),
-        }
+        true_arr = np.array(true_list)
+        est_arr = np.array(est_list)
+        lo_arr = np.array(lo_list)
+        hi_arr = np.array(hi_list)
 
-    # ---- determine subplot layout ------------------------------------------
-    subplot_specs = []  # (param_name, column_index, display_label)
-    for name in var_names:
-        n_elements = summaries[name]["true_values"].shape[1]
+        summaries = np.stack([true_arr, est_arr, lo_arr, hi_arr], axis=-1)
+        parameter_dims = true_arr.shape[1:]
+        n_elements = np.sum(parameter_dims, dtype=int)
         if n_elements == 1:
-            subplot_specs.append((name, 0, name))
+            labels = [name]
+        elif len(parameter_dims) == 1:
+            labels = [f"{name}[{i}]" for i in range(n_elements)]
         else:
-            for j in range(n_elements):
-                subplot_specs.append((name, j, f"{name}[{j}]"))
-
-    n_subplots = len(subplot_specs)
-
-    # ---- create or validate axes -------------------------------------------
-    created_figure = axes is None
-    if created_figure:
-        n_cols = min(3, n_subplots)
-        n_rows = -(-n_subplots // n_cols)  # ceil division
-        if figsize is None:
-            figsize = (5 * n_cols, 4 * n_rows)
-        fig, axes_grid = plt.subplots(n_rows, n_cols, figsize=figsize, squeeze=False)
-        axes_flat = axes_grid.flatten()
-        for k in range(n_subplots, len(axes_flat)):
-            axes_flat[k].set_visible(False)
-    else:
-        axes_flat = np.atleast_1d(axes).flatten()
-        if len(axes_flat) < n_subplots:
-            raise ValueError(f"Need {n_subplots} axes but only {len(axes_flat)} provided.")
-        fig = axes_flat[0].get_figure()
-
-    # ---- plot each subplot -------------------------------------------------
-    binomial_se = np.sqrt(ci_prob * (1 - ci_prob) / n_sims)
-
-    for i, (name, col, label) in enumerate(subplot_specs):
-        ax = axes_flat[i]
-        s = summaries[name]
-
-        true_vals = s["true_values"][:, col]
-        estimates = s["point_estimates"][:, col]
-        ci_lo = s["ci_lower"][:, col]
-        ci_hi = s["ci_upper"][:, col]
-
-        covered = (ci_lo <= true_vals) & (true_vals <= ci_hi)
-        empirical_cov = covered.mean()
-
-        yerr = np.array(
-            [
-                np.maximum(estimates - ci_lo, 0),
-                np.maximum(ci_hi - estimates, 0),
-            ]
-        )
+            labels = [f"{name}[{str(idx).strip(')(')}]" for idx in np.ndindex(parameter_dims)]
 
-        # Plot covered and missed points separately for the legend
-        mask_cov = covered
-        mask_miss = ~covered
-
-        if mask_cov.any():
-            ax.errorbar(
-                true_vals[mask_cov],
-                estimates[mask_cov],
-                yerr=yerr[:, mask_cov],
-                fmt="o",
-                color="C0",
-                markersize=4,
-                linewidth=1,
-                alpha=0.7,
-                capsize=2,
-                label="Covered",
-            )
-        if mask_miss.any():
-            ax.errorbar(
-                true_vals[mask_miss],
-                estimates[mask_miss],
-                yerr=yerr[:, mask_miss],
-                fmt="o",
-                color="C3",
-                markersize=4,
-                linewidth=1,
-                alpha=0.7,
-                capsize=2,
-                label="Missed",
-            )
-
-        # y = x reference line
-        all_vals = np.concatenate([true_vals, estimates, ci_lo, ci_hi])
-        vmin, vmax = all_vals.min(), all_vals.max()
-        margin = (vmax - vmin) * 0.05 if vmax > vmin else 0.5
-        ax.plot(
-            [vmin - margin, vmax + margin],
-            [vmin - margin, vmax + margin],
-            "k--",
-            alpha=0.4,
-            linewidth=1,
-            zorder=0,
+        param_labels.extend(labels)
+        summaries_list.append(summaries.reshape(n_sims, n_elements, 4))
+        covered_list.append(
+            ((lo_arr <= true_arr) & (true_arr <= hi_arr)).reshape(n_sims, n_elements)
         )
 
-        ax.set_xlabel("True value")
-        ax.set_ylabel(f"Posterior {point_estimate}")
-        ax.set_title(
-            f"{label} — coverage: {empirical_cov:.1%} "
-            f"(expected: {ci_prob:.1%} ± {2 * binomial_se:.1%})",
-            fontsize=9,
-            wrap=True,
-        )
-        if mask_miss.any():
-            ax.legend(fontsize="small")
+    recovery = xr.DataArray(
+        np.concatenate(summaries_list, axis=1),
+        dims=["simulation", "parameter", "quantity"],
+        coords={
+            "parameter": param_labels,
+            "quantity": ["true", "estimate", "ci_low", "ci_high"],
+        },
+    )
+    covered = xr.DataArray(
+        np.concatenate(covered_list, axis=1),
+        dims=["simulation", "parameter"],
+        coords={"parameter": param_labels},
+    )
 
-    if created_figure:
-        fig.tight_layout()
+    return xr.Dataset({"recovery": recovery, "covered": covered})
 
-    if if_show:
-        plt.show()
 
-    return fig
+def plot_parameter_recovery(
+    sbc,
+    ci_prob=None,
+    point_estimate=None,
+    var_names=None,
+    col_wrap=4,
+    figsize=None,
+    title_size="medium",
+    backend=None,
+    plot_collection=None,
+    visuals=None,
+    pc_kwargs=None,
+    if_show=True,
+):
+    """Create a parameter recovery plot for an SBC object.
 
+    For each parameter (scalar or flattened vector element), the plot shows the
+    posterior point estimate against the true value, a vertical credible interval,
+    and a 45-degree reference line.  Subplot titles report the observed coverage
+    and its expected value (plus/minus two binomial standard errors).
 
-def plot_ecdf(sbc, if_show=True):
+    Parameters
+    ----------
+    sbc : simuk.SBC
+        Fitted SBC object with posteriors and kept simulation parameters.
+    ci_prob : float, optional
+        Credible interval probability.  Defaults to ``rcParams["stats.ci_prob"]``
+        (0.89).
+    point_estimate : {"mean", "median"}, optional
+        Point estimate to plot on the y-axis.  Defaults to
+        ``rcParams["stats.point_estimate"]`` ("mean").
+    var_names : str or list of str, optional
+        Variable names to include in the plot.  If None, all parameters in
+        ``sbc.kept_simulation_params.var_names`` are used.
+    col_wrap : int, optional
+        Maximum number of columns before wrapping to the next row.  Default is 4.
+    figsize : tuple of float, optional
+        Base figure size in inches for a single subplot.  Default is (6, 6).
+    title_size : str or float, optional
+        Font size for subplot titles.  Default is "medium".
+    backend : str, optional
+        Plotting backend.  Defaults to ``rcParams["plot.backend"]``.
+    plot_collection : arviz_plots.PlotCollection, optional
+        Existing plot collection to add to.  If None, a new one is created.
+    visuals : dict, optional
+        Visuals configuration.  Keys are visual names, values are kwargs dicts.
+        Use ``False`` to disable a visual.  Default is ``{}``.
+
+        Supported visuals
+        -----------------
+        - ``"labels"``: kwargs for ``labelled_x`` / ``labelled_y``
+        - ``"reference_line"``: kwargs for the 45-degree ``dline``
+        - ``"errorbar"``: kwargs for the per-simulation scatter + CI line
+        - ``"title"``: kwargs for the subplot title
+        - ``"legend"``: kwargs for the figure-level legend
+    pc_kwargs : dict, optional
+        Additional kwargs for ``PlotCollection.wrap``.  ``figure_kwargs`` is
+        merged with the computed figsize.
+    if_show : bool, optional
+        Whether to display the figure.  Default is True.
 
-    if sbc.method == "posterior":
-        fig = plot_ecdf_pit(
-            sbc.simulations,
-            group="posterior_sbc",
-            visuals={"xlabel": False},
+    Returns
+    -------
+    arviz_plots.PlotCollection
+        The PlotCollection wrapping the recovery plot.
+    """
+    if not getattr(sbc, "keep_fits", False):
+        raise ValueError(
+            "plot_parameter_recovery requires keep_fits=True. Re-run SBC with keep_fits=True."
         )
+
+    if not sbc.posteriors:
+        raise ValueError(
+            "No posteriors found. Run sbc.run_simulations() with keep_fits=True first."
+        )
+
+    # Validate and set defaults
+    ci_prob = validate_or_use_rcparam(ci_prob, "stats.ci_prob")
+    point_estimate = validate_or_use_rcparam(point_estimate, "stats.point_estimate")
+    if figsize is None:
+        figsize = (6, 6)
+    if backend is None:
+        backend = rcParams["plot.backend"]
+    pc_kwargs = pc_kwargs or {}
+    visuals = validate_dict_argument(
+        visuals, valid_keys=["labels", "reference_line", "errorbar", "title", "legend"]
+    )
+
+    # Resolve variable names
+    all_var_names = sbc.kept_simulation_params.var_names
+    if var_names is None:
+        resolved_var_names = all_var_names
+    elif isinstance(var_names, str):
+        resolved_var_names = [var_names]
     else:
-        fig = plot_ecdf_pit(
-            sbc.simulations,
-            visuals={"xlabel": False},
+        resolved_var_names = list(var_names)
+    # Validate that all requested names exist
+    missing = set(resolved_var_names) - set(all_var_names)
+    if missing:
+        raise ValueError(
+            f"var_names contains unknown variables: {sorted(missing)}. "
+            f"Available: {list(all_var_names)}"
+        )
+
+    # Get backend module
+    plot_bknd = import_module(f"arviz_plots.backend.{backend}")
+
+    # Build recovery dataset
+    ds = _build_recovery_dataset(
+        sbc,
+        ci_prob=ci_prob,
+        point_estimate=point_estimate,
+        transform=sbc._transform,
+        var_names=resolved_var_names,
+    )
+
+    # Setup plot collection
+    n_params = len(ds.coords["parameter"])
+    n_cols = min(col_wrap, n_params)
+    n_rows = (n_params + n_cols - 1) // n_cols
+
+    # figsize is per-subplot; compute total and scale to dots
+    total_figsize = (figsize[0] * n_cols, figsize[1] * n_rows)
+    scaled_figsize = plot_bknd.scale_fig_size(
+        figsize=total_figsize,
+        figsize_units="inches",
+    )
+
+    # Merge figsize into pc_kwargs
+    pc_kwargs = pc_kwargs.copy()
+    figure_kwargs = pc_kwargs.setdefault("figure_kwargs", {})
+    figure_kwargs["figsize"] = scaled_figsize
+    figure_kwargs["figsize_units"] = "dots"
+
+    if plot_collection is None:
+        pc = PlotCollection.wrap(
+            ds, cols=["parameter"], col_wrap=col_wrap, backend=backend, **pc_kwargs
+        )
+    else:
+        pc = plot_collection
+
+    # Apply visuals
+    # Axis labels
+    visuals_labels = get_visual_kwargs(visuals, "labels", {})
+    if visuals_labels is not False:
+        pc.map(azp.visuals.labelled_x, text="True Value", **visuals_labels)
+        pc.map(azp.visuals.labelled_y, text="Posterior Mean", **visuals_labels)
+
+    # Reference line
+    visuals_ref_line = get_visual_kwargs(visuals, "reference_line", {})
+    if visuals_ref_line is not False:
+        true_vals = ds["recovery"].sel(quantity="true")
+        ci_low = ds["recovery"].sel(quantity="ci_low")
+        ci_hi = ds["recovery"].sel(quantity="ci_high")
+        line_min = np.min([true_vals.min().item(), ci_low.min().item()])
+        line_max = np.max([true_vals.max().item(), ci_hi.max().item()])
+
+        pc.map(
+            azp.visuals.dline,
+            x=[line_min, line_max],
+            y=[line_min, line_max],
+            **visuals_ref_line,
+        )
+
+    # Errorbars
+    visuals_errorbar = get_visual_kwargs(visuals, "errorbar", {})
+    if visuals_errorbar is not False:
+
+        def _recovery_errorbar(da, target, covered, **kwargs):
+            """Point estimate + vertical CI for one parameter facet."""
+            plot_backend = backend_from_object(target)
+            for sim in da.simulation.values:
+                true_v = da.sel(simulation=sim, quantity="true").item()
+                est_v = da.sel(simulation=sim, quantity="estimate").item()
+                lo_v = da.sel(simulation=sim, quantity="ci_low").item()
+                hi_v = da.sel(simulation=sim, quantity="ci_high").item()
+                cov = covered.sel(simulation=sim).item()
+
+                color = "teal" if cov else "purple"
+                plot_backend.scatter([true_v], [est_v], target, color=color, **kwargs)
+                plot_backend.line([true_v, true_v], [lo_v, hi_v], target, color=color, **kwargs)
+
+        pc.map(
+            _recovery_errorbar,
+            data=ds["recovery"],
+            covered=ds["covered"],
+            **visuals_errorbar,
+        )
+
+    # Title
+    visuals_title = get_visual_kwargs(visuals, "title", {})
+    if visuals_title is not False:
+
+        def _recovery_title(da, target, ci_prob, size, **kwargs):
+            """Add a title with coverage statistics to a parameter facet."""
+            plot_backend = backend_from_object(target)
+            label = da.coords["parameter"].item()
+            n_sims = da.sizes["simulation"]
+            coverage = da.mean().item()
+            se = np.sqrt(ci_prob * (1 - ci_prob) / n_sims)
+            margin = 2 * se
+            text = f"{label}\ncoverage: {coverage:.1%}\n(expected: {ci_prob:.1%} ± {margin:.1%})"
+            return plot_backend.title(text, target, size=size, **kwargs)
+
+        pc.map(
+            _recovery_title,
+            data=ds["covered"],
+            ci_prob=ci_prob,
+            size=title_size,
+            **visuals_title,
+        )
+
+    # Legend (figure-level)
+    visuals_legend = get_visual_kwargs(visuals, "legend", {})
+    if visuals_legend is not False:
+        plot_bknd.legend(
+            pc,
+            kwarg_list=[
+                {"color": "teal", "linestyle": "-", "width": 2},
+                {"color": "purple", "linestyle": "-", "width": 2},
+            ],
+            label_list=["Covered", "Not covered"],
+            title="Credible interval",
+            **visuals_legend,
         )
 
     if if_show:
-        plt.show()
+        pc.show()
 
-    return fig
+    return pc
diff --git a/simuk/tests/conftest.py b/simuk/tests/conftest.py
deleted file mode 100644
index 4eee837..0000000
--- a/simuk/tests/conftest.py
+++ /dev/null
@@ -1,10 +0,0 @@
-# This is needed as running all tests together now causes the test to fail
-# when the same doesn't happen when running tests separately.
-# The Error is the same: AttributeError: module 'numpyro.infer' has no
-# attribute 'initialization' in arviz_base.io_numpyro
-
-# This may have something to do with the lazy loading used in arviz_base.io_numpyro,
-# in accordance with import ordering and different behaviors when using from-imports.
-# However, I'm unable to recreate it using a minimal example. This at least enforces
-# the numpyro.infer to have the initialization attribute.
-import numpyro.infer.initialization  # noqa: F401
diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py
index 9101f94..422a679 100644
--- a/simuk/tests/test_plot.py
+++ b/simuk/tests/test_plot.py
@@ -1,5 +1,4 @@
 import arviz_plots as azp
-import matplotlib.pyplot as plt
 import numpy as np
 import numpyro
 import numpyro.distributions as dist
@@ -153,121 +152,99 @@ def test_ppr_requires_completed_simulations_posterior():
 
 def test_ppr_basic(sbc_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_with_fits, if_show=False)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
 def test_plot_ppr_basic_numpyro(sbc_with_fits_numpyro):
     fig = simuk.plot_parameter_recovery(sbc_with_fits_numpyro, if_show=False)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
 def test_ppr_basic_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, if_show=False)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
 def test_ppr_var_names_filter(sbc_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_with_fits, var_names=["mu"], if_show=False)
-    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
-    assert len(visible_axes) == 1
+
+    used_params = len(fig.data.coords["parameter"])
+    assert used_params == 1
 
 
 def test_ppr_var_names_filter_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, var_names=["mu"], if_show=False)
-    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
-    assert len(visible_axes) == 1
-
-
-def test_ppr_with_transform(sbc_with_fits):
-    fig = simuk.plot_parameter_recovery(
-        sbc_with_fits, transform=lambda name, val: np.mean(val), if_show=False
-    )
-    assert isinstance(fig, plt.Figure)
-    # The mean transform reduces theta (8,) to scalar → 3 subplots total
-    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
-    assert len(visible_axes) == 3  # mu, tau, theta (all scalar after transform)
-
-
-def test_ppr_with_bad_transform(sbc_with_fits):
-    with pytest.raises(ValueError, match="`transform` should be a function or None"):
-        simuk.plot_parameter_recovery(sbc_with_fits, transform="bad transform", if_show=False)
-
-
-def test_ppr_with_transform_posterior(sbc_posterior_with_fits):
-    fig = simuk.plot_parameter_recovery(
-        sbc_posterior_with_fits, transform=lambda name, val: np.mean(val), if_show=False
-    )
-    assert isinstance(fig, plt.Figure)
-    visible_axes = [ax for ax in fig.get_axes() if ax.get_visible()]
-    assert len(visible_axes) == 2  # mu, sigma (all scalar after transform)
+    used_params = len(fig.data.coords["parameter"])
+    assert used_params == 1
 
 
 def test_ppr_custom_ci_prob(sbc_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_with_fits, ci_prob=0.5, if_show=False)
-    plt.close(fig)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
 def test_ppr_custom_ci_prob_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, ci_prob=0.5, if_show=False)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
-def test_ppr_with_preexisting_axes(sbc_with_fits):
-    # mu(1) + tau(1) + theta(8) = 10 subplots needed
-    fig, axes = plt.subplots(2, 5)
-    returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, axes=axes, if_show=False)
-    assert returned_fig is fig
+def test_ppr_with_preexisting_plot_collection(sbc_with_fits):
+    from simuk.plots import _build_recovery_dataset  # noqa: PLC0415
 
+    ds = _build_recovery_dataset(
+        sbc_with_fits,
+        ci_prob=0.89,
+        point_estimate="mean",
+        transform=sbc_with_fits._transform,
+        var_names=sbc_with_fits.kept_simulation_params.var_names,
+    )
 
-def test_ppr_with_insufficient_preexisting_axes(sbc_with_fits):
+    pc = azp.plot_collection.PlotCollection.wrap(
+        ds,
+        cols=["parameter"],
+        col_wrap=4,
+        backend="matplotlib",
+    )
     # mu(1) + tau(1) + theta(8) = 10 subplots needed
-    with pytest.raises(ValueError, match="axes but only"):
-        _, axes = plt.subplots(2, 4)
-        simuk.plot_parameter_recovery(sbc_with_fits, axes=axes, if_show=False)
-
-
-def test_ppr_with_preexisting_axes_posterior(sbc_posterior_with_fits):
-    # mu(1) + sigma(1) + theta(8) = 10 subplots needed
-    fig, axes = plt.subplots(2, 5)
-    returned_fig = simuk.plot_parameter_recovery(sbc_posterior_with_fits, axes=axes, if_show=False)
-    assert returned_fig is fig
+    returned_fig = simuk.plot_parameter_recovery(sbc_with_fits, plot_collection=pc, if_show=False)
+    assert returned_fig is pc
 
 
 def test_ppr_median_point_estimate(sbc_with_fits):
     fig = simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="median", if_show=False)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
 def test_ppr_median_point_estimate_posterior(sbc_posterior_with_fits):
     fig = simuk.plot_parameter_recovery(
         sbc_posterior_with_fits, point_estimate="median", if_show=False
     )
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
 
 
 def test_ppr_invalid_point_estimate(sbc_with_fits):
-    with pytest.raises(ValueError, match="point_estimate"):
-        simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="mode", if_show=False)
+    with pytest.raises(ValueError, match="is not one of"):
+        simuk.plot_parameter_recovery(sbc_with_fits, point_estimate="modddd", if_show=False)
 
 
 def test_ppr_invalid_point_estimate_posterior(sbc_posterior_with_fits):
-    with pytest.raises(ValueError, match="point_estimate"):
-        simuk.plot_parameter_recovery(sbc_posterior_with_fits, point_estimate="mode", if_show=False)
+    with pytest.raises(ValueError, match="is not one of"):
+        simuk.plot_parameter_recovery(
+            sbc_posterior_with_fits, point_estimate="modddd", if_show=False
+        )
 
 
 def test_ppr_show_branch(monkeypatch, sbc_with_fits):
     called = {"value": False}
 
-    def fake_show():
+    def fake_show(*args, **kwargs):
         called["value"] = True
 
-    monkeypatch.setattr(plt, "show", fake_show)
+    monkeypatch.setattr(azp.plot_collection.PlotCollection, "show", fake_show)
 
     fig = simuk.plot_parameter_recovery(sbc_with_fits, if_show=True)
-    assert isinstance(fig, plt.Figure)
+    assert isinstance(fig, azp.plot_collection.PlotCollection)
     assert called["value"]
-    plt.close(fig)
 
 
 def test_plot_ecdf_basic(sbc_with_fits, sbc_no_fits):
@@ -289,12 +266,11 @@ def test_plot_ecdf_posterior(sbc_posterior_with_fits, sbc_posterior_no_fits):
 def test_ecdf_show_branch(monkeypatch, sbc_with_fits):
     called = {"value": False}
 
-    def fake_show():
+    def fake_show(*args, **kwargs):
         called["value"] = True
 
-    monkeypatch.setattr(plt, "show", fake_show)
+    monkeypatch.setattr(azp.plot_collection.PlotCollection, "show", fake_show)
 
     fig = simuk.plot_ecdf(sbc_with_fits, if_show=True)
     assert isinstance(fig, azp.plot_collection.PlotCollection)
     assert called["value"]
-    plt.close("all")