Skip to content

Add new tutorial on how to run pathfinder from pymc-experimental. #429

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Oct 12, 2022
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions examples/references.bib
Original file line number Diff line number Diff line change
Expand Up @@ -599,3 +599,9 @@ @article{zhang2017moderation
year = {2017},
publisher = {American Psychological Association}
}
@article{zhang2021pathfinder,
title = {Pathfinder: Parallel quasi-Newton variational inference},
author = {Zhang, Lu and Carpenter, Bob and Gelman, Andrew and Vehtari, Aki},
journal = {arXiv preprint arXiv:2108.03782},
year = {2021}
}
238 changes: 238 additions & 0 deletions examples/variational_inference/pathfinder.ipynb

Large diffs are not rendered by default.

94 changes: 94 additions & 0 deletions myst_nbs/variational_inference/pathfinder.myst.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
---
jupytext:
text_representation:
extension: .md
format_name: myst
format_version: 0.13
jupytext_version: 1.13.7
kernelspec:
display_name: pymc4
language: python
name: pymc4
---

(pathfinder)=

# Pathfinder Variational Inference

:::{post} Sept 30, 2022
:tags: variational inference, jax
:category: advanced, how-to
:author: Thomas Wiecki
:::

+++

Pathfinder {cite:p}`zhang2021pathfinder` is a variational inference algorithm that produces samples from the posterior of a Bayesian model. It compares favorably to the widely used ADVI algorithm. On large problems, it should scale better than most MCMC algorithms, including dynamic HMC (i.e. NUTS), at the cost of a more biased estimate of the posterior. For details on the algorithm, see the [arxiv preprint](https://arxiv.org/abs/2108.03782).

This algorithm is [implemented](https://github.com/blackjax-devs/blackjax/pull/194) in [BlackJAX](https://github.com/blackjax-devs/blackjax), a library of inference algorithms for [JAX](https://github.com/google/jax). Through PyMC's JAX-backend (through [aesara](https://github.com/aesara-devs/aesara)) we can run BlackJAX's pathfinder on any PyMC model with some simple wrapper code.

This wrapper code is implemented in [pymcx](https://github.com/pymc-devs/pymcx/). This tutorial shows how to run Pathfinder on your PyMC model.

You first need to install `pymcx`:

`pip install git+https://github.com/pymc-devs/pymcx`

```{code-cell} ipython3
import arviz as az
import numpy as np
import pymc as pm
import pymcx as pmx

print(f"Running on PyMC v{pm.__version__}")
```

First, define your PyMC model. Here, we use the 8-schools model.

```{code-cell} ipython3
# Data of the Eight Schools Model
J = 8
y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])

with pm.Model() as model:
mu = pm.Normal("mu", mu=0.0, sigma=10.0)
tau = pm.HalfCauchy("tau", 5.0)

theta = pm.Normal("theta", mu=0, sigma=1, shape=J)
theta_1 = mu + tau * theta
obs = pm.Normal("obs", mu=theta, sigma=sigma, shape=J, observed=y)
```

Next, we call `pmx.fit()` and pass in the algorithm we want it to use.

```{code-cell} ipython3
with model:
idata = pmx.fit(method="pathfinder")
```

Just like `pymc.sample()`, this returns an idata with samples from the posterior. Note that because these samples do not come from an MCMC chain, convergence can not be assessed in the regular way.

```{code-cell} ipython3
az.plot_trace(idata);
```

## References

:::{bibliography}
:filter: docname in docnames
:::

+++

## Authors

* Authored by Thomas Wiecki on Oct 11 2022 ([pymc-examples#429](https://github.com/pymc-devs/pymc-examples/pull/429))

+++

## Watermark

```{code-cell} ipython3
%load_ext watermark
%watermark -n -u -v -iv -w -p aesara,xarray
```