{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# A First Posterior with NLE\n", "\n", "This tutorial trains a Neural Likelihood Estimator ({term}`NLE`) on a\n", "linear-Gaussian simulator and samples the posterior with PyMC and NumPyro.\n", "The simulator is deliberately simple: its likelihood is a known Gaussian, so\n", "every step can be checked against an exact answer. The workflow is the same\n", "four phases used for any real simulator, {term}`Simulate → Train → Validate → Export`,\n", "and the core pipeline is about twenty lines of setu code.\n", "\n", "Running this notebook requires the `pymc` and `numpyro` extras (see the\n", "[install how-to](../how-to/install.md))." ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "import arviz as az\n", "import jax\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import pandas as pd\n", "import pymc as pm\n", "\n", "import setu\n", "from setu.integration.pymc import to_pymc" ] }, { "cell_type": "code", "execution_count": null, "id": "2", "metadata": {}, "outputs": [], "source": [ "# Shared setu-docs figure style: applied when this notebook is executed by the\n", "# docs build (setu_docs_style is on the build's PYTHONPATH); skipped harmlessly\n", "# when the notebook is run standalone.\n", "try:\n", " from setu_docs_style import set_style\n", "\n", " set_style()\n", "except ImportError:\n", " pass" ] }, { "cell_type": "markdown", "id": "3", "metadata": {}, "source": [ "## The simulator\n", "\n", "The simulator draws an observation from a Gaussian centered on a shifted\n", "parameter vector,\n", "\n", "$$x \\mid \\theta \\sim \\mathcal{N}(\\theta + s, \\Sigma),$$\n", "\n", "with a standard normal prior $\\theta \\sim \\mathcal{N}(0, I)$ in three\n", "dimensions. setu ships this as {func}`setu.linear_gaussian`, one of its test\n", "simulators. Because prior and likelihood are conjugate, the exact posterior is\n", "a Gaussian with a closed-form mean and covariance, which makes this the ideal\n", "first problem: the entire neural pipeline can be compared against ground\n", "truth at the end.\n", "\n", "One detail is worth flagging immediately: the parameter dimension and the\n", "observation dimension are both 3. That makes it impossible for setu to catch\n", "a swapped `theta`/`x` argument by shape, a point this tutorial returns to\n", "below." ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": {}, "outputs": [], "source": [ "theta_dim = x_dim = 3\n", "shift = jnp.array([0.5, -0.5, 1.0])\n", "likelihood_cov = 0.5 * jnp.eye(x_dim)\n", "prior_mean = jnp.zeros(theta_dim)\n", "prior_cov = jnp.eye(theta_dim)\n", "\n", "key = jax.random.key(0)\n", "k_simulate, k_create, k_fit, k_val, k_validate, k_obs = jax.random.split(key, 6)" ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "## Phase 1: Simulate\n", "\n", "{func}`setu.simulate_dataset` builds a training set from two functions: a\n", "`prior_fn` that draws a single parameter vector from a key, and a\n", "`simulator_fn` that maps one `(theta, key)` pair to one observation. Both are\n", "`jax.vmap`-ed internally over `n_samples` fresh keys.\n", "\n", "The library helpers have slightly different signatures, so two thin wrappers\n", "adapt them: {func}`setu.sample_prior` draws a whole batch (here a batch of\n", "one), and {func}`setu.linear_gaussian` takes its key as the last of four\n", "arguments. Passing everything by keyword keeps the adaptation explicit." ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "def prior(key):\n", " return setu.sample_prior(mean=prior_mean, cov=prior_cov, n_samples=1, key=key)[0]\n", "\n", "\n", "def simulator(theta, key):\n", " return setu.linear_gaussian(theta=theta, shift=shift, cov=likelihood_cov, key=key)\n", "\n", "\n", "dataset = setu.simulate_dataset(simulator, prior, n_samples=50_000, key=k_simulate)\n", "dataset.theta.shape, dataset.x.shape" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "## The observed data\n", "\n", "Inference needs something to infer from. A known ground-truth parameter\n", "$\\theta^*$ generates ten observed {term}`Trials`; the goal of the rest of the\n", "tutorial is to recover $\\theta^*$ from these ten observations alone, and the\n", "known value makes the recovery checkable." ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "theta_true = jnp.array([0.8, -0.3, 0.2])\n", "obs_keys = jax.random.split(k_obs, 10)\n", "x_obs = jax.vmap(lambda k: simulator(theta_true, k))(obs_keys)\n", "x_obs.shape" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "## Phase 2: Train\n", "\n", "setu separates model construction from training: {func}`setu.create_nle`\n", "builds an untrained estimator (by default a neural spline flow), and\n", "{func}`setu.fit_nle` trains it with early stopping on a held-out validation\n", "split. The returned {class}`setu.TrainingResult` carries the trained\n", "estimator together with the loss history." ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "nle = setu.create_nle(x_dim=x_dim, theta_dim=theta_dim, key=k_create)\n", "result = setu.fit_nle(nle, dataset, max_epochs=400, show_progress=False, key=k_fit)\n", "nle = result.nle\n", "\n", "print(\n", " f\"stopped at epoch {len(result.losses)}, best epoch {result.best_epoch}, \"\n", " f\"best validation loss {result.val_losses[result.best_epoch]:.3f}\"\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "fig, ax = plt.subplots()\n", "ax.plot(result.losses, label=\"training loss\")\n", "ax.plot(result.val_losses, label=\"validation loss\")\n", "ax.axvline(result.best_epoch, color=\"#999999\", linestyle=\"--\", label=\"best epoch\")\n", "ax.set(xlabel=\"epoch\", ylabel=\"negative log likelihood\")\n", "ax.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "12", "metadata": {}, "source": [ "## A note on argument order\n", "\n", "The trained estimator evaluates $\\log p(x \\mid \\theta)$ via\n", "`nle.log_prob(x, theta)`, with the observation first. The sampling and export\n", "functions take `theta` first. In this problem `x_dim == theta_dim`, so\n", "swapping the two arguments raises no error and returns a finite, plausible,\n", "wrong number:" ] }, { "cell_type": "code", "execution_count": null, "id": "13", "metadata": {}, "outputs": [], "source": [ "lp_correct = nle.log_prob(x=x_obs[0], theta=theta_true)\n", "lp_swapped = nle.log_prob(x=theta_true, theta=x_obs[0])\n", "\n", "print(f\"log p(x | theta) = {lp_correct:.3f}\")\n", "print(f\"swapped arguments = {lp_swapped:.3f} (finite, plausible, wrong)\")" ] }, { "cell_type": "markdown", "id": "14", "metadata": {}, "source": [ "Whenever the parameter and observation dimensions match, only the argument\n", "order distinguishes the two, and a silent swap corrupts every downstream\n", "result. Passing `x=` and `theta=` by keyword, as done throughout this\n", "tutorial, removes the failure mode entirely." ] }, { "cell_type": "markdown", "id": "15", "metadata": {}, "source": [ "## Phase 3: Validate\n", "\n", "A learned likelihood is not trusted until it is checked. {term}`Validation`\n", "runs on a fresh dataset that played no role in training. The\n", "`result.validate()` convenience runs the standard battery: training\n", "diagnostics, {term}`C2ST` (a classifier tries to tell estimator samples from\n", "simulator samples; accuracy near 0.5 means indistinguishable), MMD,\n", "sliced-Wasserstein, likelihood {term}`SBC` (a rank-calibration check on the\n", "learned density), and a likelihood-accuracy check." ] }, { "cell_type": "code", "execution_count": null, "id": "16", "metadata": {}, "outputs": [], "source": [ "val = setu.simulate_dataset(simulator, prior, n_samples=1_000, key=k_val)\n", "val_data = setu.ValidationDataset(theta=val.theta, x=val.x)\n", "\n", "validation = result.validate(val_data, key=k_validate)\n", "print(validation.summary())" ] }, { "cell_type": "markdown", "id": "17", "metadata": {}, "source": [ "Because this simulator has a tractable likelihood, a more direct check is\n", "also available here (and only here): evaluating the learned and the exact\n", "log-density on the same validation pairs. On a real problem this plot does\n", "not exist, which is exactly why the checks above matter." ] }, { "cell_type": "code", "execution_count": null, "id": "18", "metadata": {}, "outputs": [], "source": [ "exact_lp = jax.vmap(\n", " lambda x, theta: setu.analytical_log_likelihood(\n", " x=x, theta=theta, shift=shift, cov=likelihood_cov\n", " )\n", ")(val.x, val.theta)\n", "learned_lp = jax.vmap(nle.log_prob)(val.x, val.theta)\n", "\n", "fig, ax = plt.subplots(figsize=(4.5, 4.5))\n", "ax.scatter(exact_lp, learned_lp, s=8, alpha=0.4)\n", "lims = [float(exact_lp.min()), float(exact_lp.max())]\n", "ax.plot(lims, lims, color=\"#999999\", linestyle=\"--\", label=\"perfect agreement\")\n", "ax.set(xlabel=\"exact log likelihood\", ylabel=\"learned log likelihood\")\n", "ax.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "19", "metadata": {}, "source": [ "## Phase 4: Export to PyMC\n", "\n", "{term}`Export` produces a likelihood term inside an ordinary PyMC model:\n", "`to_pymc` wraps the trained estimator as a potential attached to any PyMC\n", "random variable. Priors, transformations, and sampling all remain plain PyMC.\n", "With `nuts_sampler=\"blackjax\"` the whole model graph is compiled to JAX; the\n", "default NUTS sampler works as well but calls back into Python on every step.\n", "`chain_method=\"vectorized\"` runs the chains as one vectorized computation,\n", "which is what a single device (one CPU or GPU) needs, and `random_seed` makes\n", "the run reproducible." ] }, { "cell_type": "code", "execution_count": null, "id": "20", "metadata": {}, "outputs": [], "source": [ "with pm.Model():\n", " theta = pm.Normal(\"theta\", mu=0.0, sigma=1.0, shape=theta_dim)\n", " to_pymc(nle, theta=theta, x_observed=x_obs)\n", " idata = pm.sample(\n", " draws=500,\n", " tune=500,\n", " chains=4,\n", " nuts_sampler=\"blackjax\",\n", " nuts_sampler_kwargs={\"chain_method\": \"vectorized\"},\n", " progressbar=False,\n", " random_seed=0,\n", " )\n", "\n", "az.summary(idata, var_names=[\"theta\"])" ] }, { "cell_type": "markdown", "id": "21", "metadata": {}, "source": [ "## Checking against the exact posterior\n", "\n", "With a Gaussian prior and a Gaussian likelihood, the posterior after $n$\n", "observations is Gaussian with precision and mean\n", "\n", "$$\\Lambda = \\Sigma_0^{-1} + n\\,\\Sigma^{-1}, \\qquad\n", "\\mu_{\\text{post}} = \\Lambda^{-1}\\Big(\\Sigma_0^{-1}\\mu_0\n", " + \\Sigma^{-1}\\textstyle\\sum_i (x_i - s)\\Big).$$\n", "\n", "Overlaying it on the sampled posterior closes the loop: the samples came from\n", "a neural approximation of the likelihood, yet they should trace the exact\n", "answer." ] }, { "cell_type": "code", "execution_count": null, "id": "22", "metadata": {}, "outputs": [], "source": [ "n_obs = x_obs.shape[0]\n", "post_prec = jnp.linalg.inv(prior_cov) + n_obs * jnp.linalg.inv(likelihood_cov)\n", "post_cov = jnp.linalg.inv(post_prec)\n", "post_mean = post_cov @ (\n", " jnp.linalg.inv(prior_cov) @ prior_mean\n", " + jnp.linalg.inv(likelihood_cov) @ (x_obs - shift).sum(axis=0)\n", ")\n", "\n", "pymc_samples = idata.posterior[\"theta\"].values.reshape(-1, theta_dim)\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(10, 3.2))\n", "for d, ax in enumerate(axes):\n", " ax.hist(\n", " pymc_samples[:, d],\n", " bins=40,\n", " density=True,\n", " alpha=0.6,\n", " color=\"#0072B2\",\n", " label=\"NLE + PyMC\",\n", " )\n", " sd = float(jnp.sqrt(post_cov[d, d]))\n", " grid = np.linspace(post_mean[d] - 4 * sd, post_mean[d] + 4 * sd, 200)\n", " pdf = np.exp(-0.5 * ((grid - float(post_mean[d])) / sd) ** 2)\n", " pdf /= sd * np.sqrt(2 * np.pi)\n", " ax.plot(grid, pdf, color=\"#000000\", label=\"exact posterior\")\n", " ax.axvline(\n", " float(theta_true[d]), color=\"#999999\", linestyle=\"--\", label=\"true theta\"\n", " )\n", " ax.set(title=f\"theta[{d}]\", xlabel=\"value\")\n", "axes[0].set(ylabel=\"density\")\n", "axes[0].legend(fontsize=9)\n", "fig.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "23", "metadata": {}, "source": [ "## Export to NumPyro\n", "\n", "The same trained estimator exports to NumPyro without retraining; this is\n", "what {term}`Amortized` means in practice. NumPyro is JAX-native, so\n", "`to_numpyro` adds the log-likelihood directly as a factor with no bridge\n", "layer. Reproducibility works differently from PyMC: the random state is an\n", "explicit `rng_key` passed to `mcmc.run`." ] }, { "cell_type": "code", "execution_count": null, "id": "24", "metadata": {}, "outputs": [], "source": [ "import numpyro\n", "import numpyro.distributions as dist\n", "from numpyro.infer import MCMC, NUTS\n", "\n", "from setu.integration.numpyro import to_numpyro\n", "\n", "\n", "def model():\n", " theta = numpyro.sample(\"theta\", dist.Normal(jnp.zeros(theta_dim), 1.0).to_event(1))\n", " to_numpyro(nle, theta=theta, x_observed=x_obs)\n", "\n", "\n", "mcmc = MCMC(\n", " NUTS(model),\n", " num_warmup=500,\n", " num_samples=500,\n", " num_chains=4,\n", " chain_method=\"sequential\",\n", " progress_bar=False,\n", ")\n", "mcmc.run(jax.random.key(1))\n", "numpyro_samples = np.asarray(mcmc.get_samples()[\"theta\"])" ] }, { "cell_type": "markdown", "id": "25", "metadata": {}, "source": [ "Both PPL paths agree with the closed-form posterior mean up to a small\n", "residual that reflects the approximation error of the learned likelihood\n", "(more training simulations shrink it). All three sit near the ground-truth\n", "parameter, differing from it by the posterior's own uncertainty, since ten\n", "observations only pin $\\theta$ down so far:" ] }, { "cell_type": "code", "execution_count": null, "id": "26", "metadata": {}, "outputs": [], "source": [ "pd.DataFrame(\n", " {\n", " \"exact posterior mean\": np.asarray(post_mean),\n", " \"PyMC mean\": pymc_samples.mean(axis=0),\n", " \"NumPyro mean\": numpyro_samples.mean(axis=0),\n", " \"true theta\": np.asarray(theta_true),\n", " },\n", " index=[f\"theta[{d}]\" for d in range(theta_dim)],\n", ").round(3)" ] }, { "cell_type": "markdown", "id": "27", "metadata": {}, "source": [ "## Where to go next\n", "\n", "The four phases above transfer unchanged to simulators whose likelihood is\n", "genuinely intractable; only the validation scatter plot against the exact\n", "density goes away. {doc}`Tutorial 2 <01_tadpole_binomial>` applies the\n", "workflow to real experimental data with a discrete outcome ({term}`MixedNLE`).\n", "[What is SBI?](../explanation/what-is-sbi.md) gives the conceptual background\n", "for the four phases, and the [API reference](../reference/api.md) documents\n", "every function used here." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }