tutorials/16_implemented_methods.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "85a348a5-f139-4cb0-8f82-6ebeabd2c882",
   "metadata": {},
   "source": [
    "# API of implemented methods\n",
    "\n",
    "This notebook spells out the API for all algorithms implemented in the `sbi` toolbox:\n",
    "\n",
    "- Posterior estimation (NPE)\n",
    "\n",
    "- Likelihood estimation (NLE)\n",
    "\n",
    "- Likelihood-ratio estimation (NRE)\n",
    "\n",
    "- Utilities\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e2d1852-898e-49ed-85ef-888f0abd7e4e",
   "metadata": {},
   "source": [
    "## Posterior estimation (NPE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e7559f4-d97b-4bca-8e22-6529c9c92d9c",
   "metadata": {},
   "source": [
    "**Fast ε-free Inference of Simulation Models with Bayesian Conditional Density Estimation**<br> by Papamakarios & Murray (NeurIPS 2016) <br>[[PDF]](https://papers.nips.cc/paper/6084-fast-free-inference-of-simulation-models-with-bayesian-conditional-density-estimation.pdf) [[BibTeX]](https://papers.nips.cc/paper/6084-fast-free-inference-of-simulation-models-with-bayesian-conditional-density-estimation/bibtex)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3d6a658",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example setup\n",
    "import torch\n",
    "\n",
    "from sbi.utils import BoxUniform\n",
    "\n",
    "# Define the prior\n",
    "num_dims = 2\n",
    "num_sims = 1000\n",
    "num_rounds = 2\n",
    "prior = BoxUniform(low=torch.zeros(num_dims), high=torch.ones(num_dims))\n",
    "simulator = lambda theta: theta + torch.randn_like(theta) * 0.1\n",
    "x_o = torch.tensor([0.5, 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "080081fb-1558-4756-8a8f-ff065dd1b400",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NPE_A\n",
    "\n",
    "inference = NPE_A(prior)\n",
    "proposal = prior\n",
    "for r in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    # NPE-A trains a Gaussian density estimator in all but the last round. In the last round,\n",
    "    # it trains a mixture of Gaussians, which is why we have to pass the `final_round` flag.\n",
    "    final_round = True if r == num_rounds - 1 else False\n",
    "    _ = inference.append_simulations(theta, x, proposal=proposal).train(final_round=final_round)\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f04c1c6-35b3-4a91-90c5-9fd0377eb8e2",
   "metadata": {},
   "source": [
    "**Automatic posterior transformation for likelihood-free inference**<br>by Greenberg, Nonnenmacher & Macke (ICML 2019) <br>[[PDF]](http://proceedings.mlr.press/v97/greenberg19a/greenberg19a.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "743eb04f-2cd5-4986-a33b-f2207b9cd5a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NPE\n",
    "\n",
    "inference = NPE(prior)\n",
    "proposal = prior\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x, proposal=proposal).train()\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "    proposal = posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1df5c9d",
   "metadata": {},
   "source": [
    "**BayesFlow: Learning complex stochastic models with invertible neural\n",
    "networks**<br> by Radev, S. T., Mertens, U. K., Voss, A., Ardizzone, L., & Köthe,\n",
    "U. (2020) (IEEE transactions on neural networks and learning systems 2020)<br>\n",
    "[Paper](https://ieeexplore.ieee.org/abstract/document/9298920)\n",
    "\n",
    "The density estimation part of BayesFlow is equivalent to single-round NPE. The\n",
    "additional contribution of the paper are several embedding networks for high-dimensional\n",
    "data including permutation invariant embeddings. Similar embeddings networks are\n",
    "implemented in `sbi` as well, under `sbi.neural_nets.embedding_nets`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "836f58e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Posterior estimation with BayesFlow is equivalent to single-round NPE.\n",
    "from sbi.inference import NPE\n",
    "\n",
    "inference = NPE(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior()\n",
    "samples = posterior.sample((1000,), x=x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4de2b24-94ce-4cbf-a675-b0c19b5200ca",
   "metadata": {},
   "source": [
    "**Truncated proposals for scalable and hassle-free simulation-based inference** <br> by Deistler, Goncalves & Macke (NeurIPS 2022) <br>[[Paper]](https://arxiv.org/abs/2210.04815)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae54b1a9-c3a6-4ee9-b687-bf8c046023c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NPE\n",
    "from sbi.utils import RestrictedPrior, get_density_thresholder\n",
    "\n",
    "inference = NPE(prior)\n",
    "proposal = prior\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train(force_first_round_loss=True)\n",
    "    posterior = inference.build_posterior().set_default_x(x_o)\n",
    "\n",
    "    accept_reject_fn = get_density_thresholder(posterior, quantile=1e-4)\n",
    "    proposal = RestrictedPrior(prior, accept_reject_fn, sample_with=\"rejection\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4379824-e775-46ad-946b-07cfc3ff4c43",
   "metadata": {},
   "source": [
    "**Flow Matching for Scalable Simulation-Based Inference** <br> by Dax, Wildberger, Buchholz, Green, Macke,\n",
    "Schölkopf (NeurIPS 2023) <br> [[Paper]](https://arxiv.org/abs/2305.17161)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2922328f-2d31-48c8-8ba4-0e0a40e5b308",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import FMPE\n",
    "\n",
    "inference = FMPE(prior)\n",
    "# FMPE does support multiple rounds of inference\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ad583ea-e140-4cf5-89eb-eb77292c77c3",
   "metadata": {},
   "source": [
    "**Neural posterior score estimation**<br>\n",
    "\n",
    "based on:  \n",
    "\n",
    "- **Compositional Score Modeling for Simulation-based Inference** by Geffner, T., Papamakarios, G., & Mnih, A. (ICML 2023) [[Paper]](https://proceedings.mlr.press/v202/geffner23a.html)  \n",
    "- **Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models** by Sharrock, L., Simons, J., Liu, S., & Beaumont, M. (ICML 2024) [[Paper]](https://arxiv.org/abs/2210.04872)  \n",
    "\n",
    "Note that currently only the single-round variant is implemented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1e49c3f-a16d-4e79-ad0b-2fb4cc9ce527",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NPSE\n",
    "\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "\n",
    "inference = NPSE(prior, sde_type=\"ve\")\n",
    "_ = inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d13f84e2-d35a-4f54-8cbf-0e4be1a38fb3",
   "metadata": {},
   "source": [
    "## Likelihood estimation (NLE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89a5783d-1b46-47f8-800b-3cda038ea447",
   "metadata": {},
   "source": [
    "**Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows**<br>by Papamakarios, Sterratt & Murray (AISTATS 2019) <br>[[PDF]](http://proceedings.mlr.press/v89/papamakarios19a/papamakarios19a.pdf) [[BibTeX]](https://gpapamak.github.io/bibtex/snl.bib)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4430dbe-ac60-4978-9695-d0a5b317ee57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NLE\n",
    "\n",
    "inference = NLE(prior)\n",
    "proposal = prior\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train()\n",
    "    posterior = inference.build_posterior(mcmc_method=\"slice_np_vectorized\",\n",
    "                                          mcmc_parameters={\"num_chains\": 20,\n",
    "                                                           \"thin\": 5})\n",
    "    proposal = posterior.set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d22866-e095-4f4e-bcef-40bc196a8703",
   "metadata": {},
   "source": [
    "**Variational methods for simulation-based inference** <br> by Glöckler, Deistler, Macke (ICLR 2022) <br>[[Paper]](https://arxiv.org/abs/2203.04176)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d284d6c5-e6f6-4b1d-9c15-d6fa1736a10e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NLE\n",
    "\n",
    "inference = NLE(prior)\n",
    "proposal = prior\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train()\n",
    "    posterior = inference.build_posterior(sample_with=\"vi\",\n",
    "                                          vi_method=\"fKL\").set_default_x(x_o)\n",
    "    proposal = posterior.train()  # Train VI posterior on given x_o."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1ca3f25-c7d2-4688-b5c0-f50864d650ba",
   "metadata": {},
   "source": [
    "**Flexible and efficient simulation-based inference for models of decision-making** <br> by Boelts, Lueckmann, Gao, Macke (Elife 2022) <br>[[Paper]](https://elifesciences.org/articles/77220)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e6d6d8f-8718-44cd-bdf7-a4af2887fc14",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import MNLE\n",
    "\n",
    "inference = MNLE(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d81ff49e-f363-43c0-ba16-3ee6a697be5e",
   "metadata": {},
   "source": [
    "## Likelihood-ratio estimation (NRE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0823041a-b3f7-4cd3-8a48-497450f622ea",
   "metadata": {},
   "source": [
    "**Likelihood-free MCMC with Amortized Approximate Likelihood Ratios**<br>by Hermans, Begy & Louppe (ICML 2020) <br>[[PDF]](http://proceedings.mlr.press/v119/hermans20a/hermans20a.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b58c3609-7bd7-40ce-a154-f72a190da2ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NRE_A\n",
    "\n",
    "inference = NRE_A(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42b12584-63cd-47d4-93d2-430d407e1e0b",
   "metadata": {},
   "source": [
    "**On Contrastive Learning for Likelihood-free Inference**<br>Durkan, Murray & Papamakarios (ICML 2020) <br>[[PDF]](http://proceedings.mlr.press/v119/durkan20a/durkan20a.pdf).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e36ab4e7-713f-4ff2-b467-8b481a149861",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NRE\n",
    "\n",
    "inference = NRE(prior)\n",
    "proposal = prior\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    _ = inference.append_simulations(theta, x).train()\n",
    "    posterior = inference.build_posterior(mcmc_method=\"slice_np_vectorized\",\n",
    "                                          mcmc_parameters={\"num_chains\": 20,\n",
    "                                                           \"thin\": 5})\n",
    "    proposal = posterior.set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d0151a",
   "metadata": {},
   "source": [
    "**Towards Reliable Simulation-Based Inference with Balanced Neural Ratio Estimation**<br>by Delaunoy, Hermans, Rozet, Wehenkel & Louppe (NeurIPS 2022) <br>[[PDF]](https://arxiv.org/pdf/2208.13624.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85e6cf8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import BNRE\n",
    "\n",
    "inference = BNRE(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train(regularization_strength=100.)\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7ecf9a0-6a82-42e1-9c12-97d0f5d825be",
   "metadata": {},
   "source": [
    "**Contrastive Neural Ratio Estimation**<br>Benjamin Kurt Miller, Christoph Weniger, Patrick Forré (NeurIPS 2022) <br>[[PDF]](https://arxiv.org/pdf/2210.06170.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ec55e76-dd86-46d1-a7cc-643324488820",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The main feature of NRE-C is producing an exact ratio of densities at optimum,\n",
    "# even when using multiple contrastive pairs (classes).\n",
    "\n",
    "from sbi.inference import NRE_C\n",
    "\n",
    "inference = NRE_C(prior)\n",
    "proposal = prior\n",
    "theta = proposal.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "_ = inference.append_simulations(theta, x).train(\n",
    "    num_classes=5,  # sees `2 * num_classes - 1` marginally drawn contrastive pairs.\n",
    "    gamma=1.0,  # controls the weight between terms in its loss function.\n",
    ")\n",
    "posterior = inference.build_posterior().set_default_x(x_o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6271d3b2-1d64-45b8-93b7-b640ab7dafc5",
   "metadata": {},
   "source": [
    "## Diagnostics and utilities\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07414c0c-f8b5-45bd-9597-0eaf19d50a13",
   "metadata": {},
   "source": [
    "**Simulation-based calibration**<br>by Talts, Betancourt, Simpson, Vehtari, Gelman (arxiv 2018) <br>[[Paper]](https://arxiv.org/abs/1804.06788))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7066ef9b-0e3d-44d3-a80e-5e06de7845ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.diagnostics import run_sbc\n",
    "from sbi.analysis import sbc_rank_plot\n",
    "\n",
    "thetas = prior.sample((1000,))\n",
    "xs = simulator(thetas)\n",
    "\n",
    "# SBC is fast for fully amortized NPE.\n",
    "inference = NPE(prior)\n",
    "theta = prior.sample((num_sims,))\n",
    "x = simulator(theta)\n",
    "inference.append_simulations(theta, x).train()\n",
    "posterior = inference.build_posterior()\n",
    "\n",
    "ranks, dap_samples = run_sbc(\n",
    "    thetas, xs, posterior, num_posterior_samples=1_000\n",
    ")\n",
    "\n",
    "fig, axes = sbc_rank_plot(\n",
    "    ranks=ranks,\n",
    "    num_posterior_samples=1000,\n",
    "    plot_type=\"hist\",\n",
    "    num_bins=20,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48853668-7b6f-4cfd-9d93-7d62f0e77de8",
   "metadata": {},
   "source": [
    "**Expected coverage (sample-based)**<br>as computed in Deistler, Goncalves, Macke (Neurips 2022) [[Paper]](https://arxiv.org/abs/2210.04815) and in Rozet, Louppe (2021) [[Paper]](https://matheo.uliege.be/handle/2268.2/12993)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e3d581-8a7f-4133-8756-9750f0174c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "thetas = prior.sample((100,))\n",
    "xs = simulator(thetas)\n",
    "\n",
    "ranks, dap_samples = run_sbc(\n",
    "    thetas,\n",
    "    xs,\n",
    "    posterior,\n",
    "    num_posterior_samples=1_000,\n",
    "    reduce_fns=posterior.log_prob  # Difference to SBC.\n",
    ")\n",
    "\n",
    "# NOTE: Here we obtain a single rank plot because ranks are calculated\n",
    "# for the entire posterior and not for each marginal like in SBC.\n",
    "fig, axes = sbc_rank_plot(\n",
    "    ranks=ranks,\n",
    "    num_posterior_samples=1000,\n",
    "    plot_type=\"hist\",\n",
    "    num_bins=20,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3962b175",
   "metadata": {},
   "source": [
    "**TARP: Sampling-Based Accuracy Testing of Posterior Estimators for General Inference**\n",
    "\n",
    "Lemos, Coogan, Hezaveh & Perreault-Levasseur (ICML 2023)<br>[[Paper]](https://arxiv.org/abs/2302.03026)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de26848",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.diagnostics.tarp import run_tarp\n",
    "from sbi.analysis import plot_tarp\n",
    "\n",
    "thetas = prior.sample((1000,))\n",
    "xs = simulator(thetas)\n",
    "\n",
    "expected_coverage, ideal_coverage = run_tarp(\n",
    "    thetas,\n",
    "    xs,\n",
    "    posterior,\n",
    "    references=None,  # optional, defaults to uniform samples across parameter space.\n",
    "    num_posterior_samples=1_000,\n",
    ")\n",
    "\n",
    "fix, axes = plot_tarp(expected_coverage, ideal_coverage)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54a88026",
   "metadata": {},
   "source": [
    "**Restriction estimator**<br>by Deistler, Macke & Goncalves (PNAS 2022) <br>[[Paper]](https://www.pnas.org/doi/10.1073/pnas.2207632119)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc5e4c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sbi.inference import NPE\n",
    "from sbi.utils import RestrictionEstimator\n",
    "\n",
    "restriction_estimator = RestrictionEstimator(prior=prior)\n",
    "proposal = prior\n",
    "\n",
    "for _ in range(num_rounds):\n",
    "    theta = proposal.sample((num_sims,))\n",
    "    x = simulator(theta)\n",
    "    restriction_estimator.append_simulations(theta, x)\n",
    "    classifier = restriction_estimator.train()\n",
    "    proposal = restriction_estimator.restrict_prior()\n",
    "\n",
    "all_theta, all_x, _ = restriction_estimator.get_simulations()\n",
    "\n",
    "inference = NPE(prior)\n",
    "density_estimator = inference.append_simulations(all_theta, all_x).train()\n",
    "posterior = inference.build_posterior()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "c50aa3a452b5e33eec699c3d0adceaddf116b15627c63bb6b43782d4547b8f5a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}