Feat parameter recovery plot

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)"
    ]
   },
   {

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)"
    ]
   },
   {

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>`.
 

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"

simuk/numpyro_adapter.py → 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

simuk/pymc_adapter.py → 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):