Merge remote-tracking branch 'origin/main' into merge_flow_builders_to_current_main

Author: manuelgloeckler

docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb

@@ -62,6 +62,7 @@
     "\n",
     "from sbi.analysis import pairplot\n",
     "from sbi.inference import NLE, NPE, simulate_for_sbi\n",
+    "from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters\n",
     "from sbi.simulators.linear_gaussian import (\n",
     "    linear_gaussian,\n",
     "    true_posterior_linear_gaussian_mvn_prior,\n",
@@ -220,17 +221,16 @@
     "nle_samples = []\n",
     "num_samples = 1000\n",
     "\n",
-    "mcmc_parameters = dict(\n",
+    "mcmc_parameters = MCMCPosteriorParameters(\n",
+    "    method=\"slice_np_vectorized\",\n",
     "    num_chains=50,\n",
     "    thin=5,\n",
     "    warmup_steps=30,\n",
     "    init_strategy=\"proposal\",\n",
     ")\n",
-    "mcmc_method = \"slice_np_vectorized\"\n",
     "\n",
     "posterior = inferer.build_posterior(\n",
-    "    mcmc_method=mcmc_method,\n",
-    "    mcmc_parameters=mcmc_parameters,\n",
+    "    posterior_parameters=mcmc_parameters,\n",
     ")\n",
     "\n",
     "# Generate samples with MCMC given the same set of x_os as above.\n",

docs/advanced_tutorials/13_diagnostics_lc2st.ipynb

@@ -240,7 +240,7 @@
     "# sample calibration data\n",
     "theta_cal = prior.sample((NUM_CAL,))\n",
     "x_cal = simulator(theta_cal)\n",
-    "post_samples_cal = posterior.sample_batched((1,), x=x_cal)[0]"
+    "post_samples_cal = posterior.sample_batched((1,), x=x_cal, max_sampling_batch_size=10)[0]"
    ]
   },
   {
@@ -881,7 +881,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.9"
   },
   "toc": {
    "base_numbering": 1,

docs/how_to_guide/10_refine_posterior_with_importance_sampling.ipynb

@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "541b890a",
+   "id": "ebc792ba",
    "metadata": {},
    "source": [
     "# How to refine posterior estimates with importance sampling"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "966b0608",
+   "id": "f91df50c",
    "metadata": {},
    "source": [
     "The `sbi` toolbox does not require that the simulator can **evaluate** the likelihood (it only requires samples). If the likelihood can be evaluated, then one can refine the posterior estimate _after training_ with likelihood evaluations. This is typically done with importance sampling (using the posterior estimate as proposal).\n",
@@ -20,15 +20,15 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0b3b8ade",
+   "id": "ea3b6412",
    "metadata": {},
    "source": [
     "## Main syntax"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "779b28e5",
+   "id": "2abe3610",
    "metadata": {},
    "source": [
     "```python\n",
@@ -54,15 +54,15 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3907af65",
+   "id": "9c480c4f",
    "metadata": {},
    "source": [
     "## Example"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "4e5c35cc",
+   "id": "a10d9afc",
    "metadata": {},
    "source": [
     "More details can be found in the [tutorial on importance sampling for refining the posterior estimate](https://sbi.readthedocs.io/en/latest/advanced_tutorials/15_importance_sampled_posteriors.html)."
@@ -85,7 +85,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,

docs/how_to_guide/11_iid_sampling_with_nle_or_nre.ipynb

@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "fa7814eb",
+   "id": "bcc62481",
    "metadata": {},
    "source": [
     "# How to sample the posterior given iid observations (for NLE and NRE)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "dd74fc35",
+   "id": "8431dd64",
    "metadata": {},
    "source": [
     "> Important: This how-to guide is only relevant if you have iid observations **and** you are using NLE or NRE. If you have iid observations and want to use NPE, you should construct a permutation-invariant embedding net, which is explained in [this how-to guide](https://sbi.readthedocs.io/en/latest/how_to_guide/08_permutation_invariant_embeddings.html).\n",
@@ -20,7 +20,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "adfe175c",
+   "id": "8366c969",
    "metadata": {},
    "source": [
     "```python\n",
@@ -40,15 +40,15 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e0b37ac2",
+   "id": "8252b623",
    "metadata": {},
    "source": [
     "## Example and further explanation"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "980ece2c",
+   "id": "4d8e197d",
    "metadata": {},
    "source": [
     "For more details, see [this tutorial](https://sbi.readthedocs.io/en/latest/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.html)."
@@ -71,7 +71,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,

docs/how_to_guide/13_diagnostics_lc2st.ipynb

@@ -24,18 +24,18 @@
    "source": [
     "```python\n",
     "from sbi.diagnostics.lc2st import LC2ST\n",
-    "from sbi.analysis.plot import pp_plot_lc2st\n",
     "\n",
-    "# Sample calibration data.\n",
-    "num_lc2st_samples = 1_000\n",
+    "# Simulate calibration data. Should be at least in the thousands.\n",
+    "num_lc2st_samples = 10_000\n",
     "prior_samples = prior.sample((num_lc2st_samples,))\n",
     "prior_predictives = simulator(prior_samples)\n",
     "\n",
     "# Generate one posterior sample for every prior predictive.\n",
-    "post_samples_cal = []\n",
-    "for x in prior_predictives:\n",
-    "    post_samples_cal.append(posterior.sample((1,), x=x)[0])\n",
-    "post_samples_cal = torch.stack(post_samples_cal)\n",
+    "post_samples_cal = posterior.sample_batched(\n",
+    "    (1,),\n",
+    "    x=prior_predictives,\n",
+    "    max_sampling_batch_size=10\n",
+    ")[0]\n",
     "\n",
     "# Train the L-C2ST classifier.\n",
     "lc2st = LC2ST(\n",
@@ -50,28 +50,75 @@
     "\n",
     "# Note: x_o must have a batch-dimension. I.e. `x_o.shape == (1, observation_shape)`.\n",
     "post_samples_star = posterior.sample((10_000,), x=x_o)\n",
-    "probs_data, _ = lc2st.get_scores(\n",
+    "probs_data, scores_data = lc2st.get_scores(\n",
     "    theta_o=post_samples_star,\n",
     "    x_o=x_o,\n",
     "    return_probs=True,\n",
     "    trained_clfs=lc2st.trained_clfs\n",
     ")\n",
-    "probs_null, _ = lc2st.get_statistics_under_null_hypothesis(\n",
+    "probs_null, scores_null = lc2st.get_statistics_under_null_hypothesis(\n",
     "    theta_o=post_samples_star,\n",
     "    x_o=x_o,\n",
-    "    return_probs=True\n",
+    "    return_probs=True,\n",
     ")\n",
     "\n",
+    "conf_alpha = 0.05\n",
+    "p_value = lc2st.p_value(post_samples_star, torch.as_tensor(x_o).unsqueeze(0))\n",
+    "reject = lc2st.reject_test(post_samples_star, torch.as_tensor(x_o).unsqueeze(0), alpha=conf_alpha)\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(5, 3))\n",
+    "quantiles = np.quantile(scores_null, [0, 1-conf_alpha])\n",
+    "ax.hist(scores_null, bins=50, density=True, alpha=0.5, label=\"Null\")\n",
+    "ax.axvline(scores_data, color=\"red\", label=\"Observed\")\n",
+    "ax.axvline(quantiles[0], color=\"black\", linestyle=\"--\", label=\"95% CI\")\n",
+    "ax.axvline(quantiles[1], color=\"black\", linestyle=\"--\")\n",
+    "ax.set_xlabel(\"Test statistic\")\n",
+    "ax.set_ylabel(\"Density\")\n",
+    "ax.set_title(f\"p-value = {p_value:.3f}, reject = {reject}\")\n",
+    "plt.show()\n",
+    "```\n",
+    "\n",
+    "This will return a figure as the following:"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "<img src=\"data/lc2st_hist_of_statistics.png\" width=\"500\">"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If the red line is outside of the two dotted black lines (as above), then L-C2ST rejects the null-hypothesis that the approximate posterior matches the true posterior (i.e., your posterior is likely wrong).\n",
+    "\n",
+    "If the posterior is wrong, then you can get insights into whether the posterior is under- or over-confident as follows:"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```python\n",
+    "from sbi.analysis.plot import pp_plot_lc2st\n",
+    "\n",
     "pp_plot_lc2st(\n",
     "    probs=[probs_data],\n",
     "    probs_null=probs_null,\n",
     "    conf_alpha=0.05,\n",
     "    labels=[\"Classifier probabilities \\n on observed data\"],\n",
     "    colors=[\"red\"],\n",
     ")\n",
-    "```\n",
-    "\n",
-    "This will return a figure as the following:"
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This plots something like the following:"
    ]
   },
   {
@@ -85,7 +132,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "If the red line lies within the gray region, the we cannot reject the null-hypothesis that the approximate posterior matches the true posterior. If the red line is below the gray area, then the `posterior` is over-confident. If the red line is above the gray area, then the `posterior` is under-confident."
+    "If the red line is below the gray area, then the `posterior` is over-confident. If the red line is above the gray area, then the `posterior` is under-confident."
    ]
   },
   {
@@ -108,7 +155,7 @@
    "source": [
     "## Citation\n",
     "\n",
-    "```\n",
+    "```text\n",
     "@article{linhart2023c2st,\n",
     "  title={L-c2st: Local diagnostics for posterior approximations in simulation-based inference},\n",
     "  author={Linhart, Julia and Gramfort, Alexandre and Rodrigues, Pedro},\n",
@@ -137,7 +184,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.9"
   },
   "toc": {
    "base_numbering": 1,

docs/how_to_guide/14_choose_diagnostic_tool.ipynb

@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "82dea70e",
+   "id": "4d43f851",
    "metadata": {},
    "source": [
     "# How to choose a diagnostic tool"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "1c744d13",
+   "id": "6abcd346",
    "metadata": {},
    "source": [
     "`sbi` implements a range of diagnostic tools. Here, we provide an overview of what they do:\n",
@@ -24,7 +24,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3660182d",
+   "id": "df220eed",
    "metadata": {},
    "source": [
     "All of these checks can provide a different angle on the correctness of the posterior, but they all require additional simulations. As such, we recommend running as many diagnostic tools as possible with your simulation budget."
@@ -47,7 +47,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,