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 `. 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/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 86% rename from simuk/numpyro_adapter.py rename to simuk/adapters/numpyro_adapter.py index 20ca7a5..fe45587 100644 --- a/simuk/numpyro_adapter.py +++ b/simuk/adapters/numpyro_adapter.py @@ -4,11 +4,12 @@ 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 -from simuk.backend_adapter import BackendAdapter +from simuk.adapters.backend_adapter import BackendAdapter log = logging.getLogger(__name__) @@ -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,15 @@ 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 +134,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 +162,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/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/plots.py b/simuk/plots.py new file mode 100644 index 0000000..6480756 --- /dev/null +++ b/simuk/plots.py @@ -0,0 +1,351 @@ +"""Parameter recovery plotting for simulation-based calibration.""" + +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 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: + fig = plot_ecdf_pit( + sbc.simulations, + group="prior_sbc", + visuals={"xlabel": False}, + ) + + if if_show: + fig.show() + + return fig + + +def _build_recovery_dataset(sbc, ci_prob, point_estimate, transform, var_names): + """Build an xarray.Dataset suitable for arviz_plots faceting. + + 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. + + Using a flat ``parameter`` dimension avoids naming conflicts between + data variables and their dimension coordinates when parameters are + scalar. + """ + ref_params = sbc.kept_simulation_params.ref_params + n_sims = len(sbc.posteriors) + alpha_lo = (1 - ci_prob) / 2 + alpha_hi = 1 - alpha_lo + + 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 = 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: + 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 = 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)) + est_list.append(np.atleast_1d(est)) + lo_list.append(np.atleast_1d(ci_lo)) + hi_list.append(np.atleast_1d(ci_hi)) + + true_arr = np.array(true_list) + est_arr = np.array(est_list) + lo_arr = np.array(lo_list) + hi_arr = np.array(hi_list) + + 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: + labels = [name] + elif len(parameter_dims) == 1: + labels = [f"{name}[{i}]" for i in range(n_elements)] + else: + labels = [f"{name}[{str(idx).strip(')(')}]" for idx in np.ndindex(parameter_dims)] + + 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) + ) + + 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}, + ) + + return xr.Dataset({"recovery": recovery, "covered": covered}) + + +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). + + 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. + + 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: + 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: + pc.show() + + return pc diff --git a/simuk/sbc.py b/simuk/sbc.py index 985c8c3..575c345 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 @@ -218,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() @@ -235,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 @@ -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): diff --git a/simuk/tests/test_plot.py b/simuk/tests/test_plot.py new file mode 100644 index 0000000..422a679 --- /dev/null +++ b/simuk/tests/test_plot.py @@ -0,0 +1,276 @@ +import arviz_plots as azp +import numpy as np +import numpyro +import numpyro.distributions as dist +import pymc as pm +import pytest +from numpyro.infer import NUTS + +import simuk + +# Test data (same as test_prior_sbc.py) +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_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") +def sbc_with_fits(): + sbc = simuk.SBC( + centered_eight, + num_simulations=10, + sample_kwargs={"draws": 10, "tune": 10}, + seed=42, + ) + sbc.run_simulations() + 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( + centered_eight, + num_simulations=10, + sample_kwargs={"draws": 10, "tune": 10}, + keep_fits=False, + seed=42, + ) + 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_pymc = pm.HalfNormal("sigma", sigma=2) + y_data = pm.Data("y_data", obs_data) + pm.Normal("y", mu=mu, sigma=sigma_pymc, 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, + seed=42, + 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}, + seed=42, + keep_fits=False, + ) + sbc.run_simulations() + return sbc + + +def test_ppr_requires_keep_fits(sbc_no_fits): + with pytest.raises(ValueError, match="keep_fits"): + 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"): + 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"): + 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"): + 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, 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, 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, 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) + + 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) + 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) + 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, azp.plot_collection.PlotCollection) + + +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, + ) + + pc = azp.plot_collection.PlotCollection.wrap( + ds, + cols=["parameter"], + col_wrap=4, + backend="matplotlib", + ) + # mu(1) + tau(1) + theta(8) = 10 subplots needed + 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, 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, azp.plot_collection.PlotCollection) + + +def test_ppr_invalid_point_estimate(sbc_with_fits): + 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="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(*args, **kwargs): + called["value"] = True + + monkeypatch.setattr(azp.plot_collection.PlotCollection, "show", fake_show) + + fig = simuk.plot_parameter_recovery(sbc_with_fits, if_show=True) + assert isinstance(fig, azp.plot_collection.PlotCollection) + assert called["value"] + + +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) + + +def test_ecdf_show_branch(monkeypatch, sbc_with_fits): + called = {"value": False} + + def fake_show(*args, **kwargs): + called["value"] = True + + 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"]