Skip to content
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

External samplers docs #403

Merged
merged 94 commits into from
Oct 19, 2023
Merged
Changes from all commits
Commits
Show all changes
94 commits
Select commit Hold shift + click to select a range
4f49a3e
Add all high level docs from Turing.jl
yebai Jan 6, 2023
912b1bc
Fix some permalinks.
yebai Jan 6, 2023
b9a8143
Converted Quickstart doc
cgbotta Jan 24, 2023
6a82b20
Converted contributing pages to jmd
cgbotta Jan 25, 2023
306f777
add project files
yebai Jan 26, 2023
892c301
Disable evaluating style examples
yebai Jan 26, 2023
712746e
Update and rename index.md to index.jmd
yebai Jan 26, 2023
1c72459
Rename index.jmd to using-turing.jmd
yebai Jan 26, 2023
8ccdb21
fix formatting for guide.md
yebai Jan 26, 2023
0ad33c2
Fixed formatting.
yebai Jan 26, 2023
2d1e31e
Update tutorials/for-developers-interface/interface.md
yebai Jan 26, 2023
ef95a28
Update tutorials/for-developers-abstractmcmc-turing/how_turing_implem…
yebai Jan 26, 2023
c70aece
Update tutorials/for-developers-abstractmcmc-turing/how_turing_implem…
yebai Jan 26, 2023
9e8c16d
update sample viz
yebai Jan 26, 2023
ff63088
add dummy julia code block to enable jmd building
yebai Jan 26, 2023
cb299b1
Bugfix for negative variance.
yebai Jan 26, 2023
bb392ad
More conversions
cgbotta Feb 2, 2023
a1879f6
Added dependencies and changed directory names
cgbotta Feb 8, 2023
341b21e
Fixed dependencies
cgbotta Apr 17, 2023
57d7894
Fixing dependencies
cgbotta Apr 17, 2023
c2efc56
Fixing dependencies
cgbotta Apr 17, 2023
af861a4
Fixing dependencies
cgbotta Apr 17, 2023
06b1f1a
Fixing dependencies
cgbotta Apr 17, 2023
f6d8ed9
Fixing dependencies
cgbotta Apr 17, 2023
98a4eb8
Fixing dependencies
cgbotta Apr 17, 2023
ee0c267
Fixing dependencies
cgbotta Apr 17, 2023
1d14387
Fixing dependencies
cgbotta Apr 17, 2023
80c455a
Fixing dependencies
cgbotta Apr 17, 2023
a4f4659
Fixing dependencies
cgbotta Apr 17, 2023
6aa8085
Fixing dependencies
cgbotta Apr 17, 2023
8f43c61
Fixing dependencies
cgbotta Apr 17, 2023
6e3433a
Fixing dependencies
cgbotta Apr 17, 2023
4cad629
Fixing dependencies
cgbotta Apr 17, 2023
051da51
Fixing dependencies
cgbotta Apr 17, 2023
2d1583f
other tutorials
cgbotta Apr 17, 2023
b47ab65
other tutorials
cgbotta Apr 17, 2023
96ecc01
other tutorials
cgbotta Apr 17, 2023
49dbdc5
Merge branch 'master' into hg/transfer-docs
yebai Jun 11, 2023
da8bbee
fix merge conflict with master
yebai Jul 6, 2023
41c8c1a
some edits to advanced
yebai Jul 6, 2023
b86ec7f
remove weave from deps
yebai Jul 6, 2023
8787423
improve compiler
yebai Jul 6, 2023
1ef0797
added julia formatting for a block which it was missing (#398)
torfjelde Jul 6, 2023
4ec62cb
AbstractMCMC fixes, rename some files, and update manifests (#399)
willtebbutt Jul 6, 2023
2f51b7b
no empty cell + hyperlink to bishop (#400)
yebai Jul 6, 2023
886d617
add numeric prefix to docs
yebai Jul 6, 2023
56f50cb
triger rebuild
yebai Jul 6, 2023
b3c814a
add Distributions to deps
yebai Jul 6, 2023
8c07963
minor fixes
yebai Jul 6, 2023
a0889fc
Minor edits
yebai Jul 6, 2023
3f0d51e
Update tutorials/docs-12-using-turing-guide/using-turing-guide.jmd
yebai Jul 6, 2023
919ba2f
Update tutorials/docs-09-using-turing-advanced/advanced.jmd
yebai Jul 6, 2023
51afa20
docs draft
JaimeRZP Jul 7, 2023
e2e1d1a
Apply suggestions from code review
JaimeRZP Jul 24, 2023
37c5863
pull
JaimeRZP Jul 24, 2023
b491bd5
upd branch
JaimeRZP Jul 24, 2023
db11514
env
JaimeRZP Jul 24, 2023
fb7fa6f
format
JaimeRZP Jul 24, 2023
a4a079b
better name
JaimeRZP Jul 24, 2023
45b96ce
undoing rename
JaimeRZP Jul 24, 2023
ad9fea4
undoing fck up
JaimeRZP Jul 24, 2023
dc5387a
working on Pathfinder example
JaimeRZP Jul 24, 2023
46773bd
Apply suggestions from code review
JaimeRZP Jul 25, 2023
150091b
citations
JaimeRZP Jul 25, 2023
2fbb15e
no propaganda
JaimeRZP Jul 25, 2023
3d0ccf4
flesing out the examples
JaimeRZP Jul 25, 2023
7eeaa20
format
JaimeRZP Jul 25, 2023
34a6ede
more concise
JaimeRZP Jul 25, 2023
a1680a3
format
JaimeRZP Jul 25, 2023
ee426db
advancedMH
JaimeRZP Jul 25, 2023
207d68b
for devs
JaimeRZP Jul 26, 2023
1af8100
Apply suggestions from code review
JaimeRZP Jul 27, 2023
1f1cb9d
Update tutorials/15-external-samplers/external-samplers.jmd
JaimeRZP Jul 27, 2023
2761834
bug
JaimeRZP Jul 27, 2023
5df99c3
Merge branch 'external_samplers_docs' of https://github.com/TuringLan…
JaimeRZP Jul 27, 2023
1b39ea7
Apply suggestions from code review
JaimeRZP Jul 28, 2023
5ed43db
upd
JaimeRZP Jul 28, 2023
4d13403
Update tutorials/15-external-samplers/external-samplers.jmd
JaimeRZP Jul 28, 2023
7796e11
Update tutorials/15-external-samplers/external-samplers.jmd
JaimeRZP Jul 28, 2023
24f1cc3
Update external-samplers.jmd
yebai Jul 30, 2023
6ec06d3
update manifest.toml
yebai Jul 30, 2023
f8e9b03
Update external-samplers.jmd
yebai Aug 17, 2023
7cc1703
update manifest file
yebai Aug 17, 2023
f79aa96
bug
JaimeRZP Sep 19, 2023
32284d5
upd deps
JaimeRZP Sep 20, 2023
1009fe5
Merge branch 'master' into external_samplers_docs
JaimeRZP Sep 20, 2023
cc451bb
Revert "upd deps"
JaimeRZP Sep 20, 2023
1b92ee7
Merge branch 'external_samplers_docs' of https://github.com/TuringLan…
JaimeRZP Sep 20, 2023
ef2b2eb
bump manifesto
JaimeRZP Sep 20, 2023
1700ebc
nearly there
JaimeRZP Sep 22, 2023
48fb69c
bump turing
JaimeRZP Oct 19, 2023
da45d17
Merge branch 'master' into external_samplers_docs
JaimeRZP Oct 19, 2023
2b48212
slight formatting error
JaimeRZP Oct 19, 2023
bae440f
Merge branch 'external_samplers_docs' of https://github.com/TuringLan…
JaimeRZP Oct 19, 2023
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
1,466 changes: 1,466 additions & 0 deletions tutorials/docs-16-external-samplers/Manifest.toml

Large diffs are not rendered by default.

8 changes: 8 additions & 0 deletions tutorials/docs-16-external-samplers/Project.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
[deps]
AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
MicroCanonicalHMC = "234d2aa0-2291-45f7-9047-6fa6f316b0a8"
Pathfinder = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
169 changes: 169 additions & 0 deletions tutorials/docs-16-external-samplers/external-samplers.jmd
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
---
title: Using External Sampler
permalink: /tutorials/:name/
mathjax: true
weave_options:
error : false
---

# Using External Samplers

`Turing` provides several wrapped samplers from external sampling libraries, e.g., HMC samplers from `AdvancedHMC`.
These wrappers allow new users to seamlessly sample statistical models without leaving `Turing`
However, these wrappers might only sometimes be complete, missing some functionality from the wrapped sampling library.
Moreover, users might want to use samplers currently not wrapped within `Turing`.

For these reasons, `Turing` also makes running external samplers on Turing models easy without any necessary modifications or wrapping!
Throughout, we will use a 10-dimensional Neal's funnel as a running example::

```julia
# Import libraries.
using Turing, Random, LinearAlgebra

d = 10
@model function funnel()
θ ~ Truncated(Normal(0, 3), -3, 3)
z ~ MvNormal(zeros(d - 1), exp(θ) * I)
return x ~ MvNormal(z, I)
end
```

Now we sample the model to generate some observations, which we can then condition on.

```julia
(; x) = rand(funnel() | (θ=0,))
model = funnel() | (; x);
```

Users can use any sampler algorithm to sample this model if it follows the `AbstractMCMC` API.
Before discussing how this is done in practice, giving a high-level description of the process is interesting.
Imagine that we created an instance of an external sampler that we will call `spl` such that `typeof(spl)<:AbstractMCMC.AbstractSampler`.
In order to avoid type ambiguity within Turing, at the moment it is necessary to declare `spl` as an external sampler to Turing `espl = externalsampler(spl)`, where `externalsampler(s::AbstractMCMC.AbstractSampler)` is a Turing function that types our external sampler adequately.

An excellent point to start to show how this is done in practice is by looking at the sampling library `AdvancedMH` ((`AdvancedMH`'s GitHub)[[https://github.com/TuringLang/AdvancedMH.jl]) for Metropolis-Hastings (MH) methods.
Let's say we want to use a random walk Metropolis-Hastings sampler without specifying the proposal distributions.
The code below constructs an MH sampler using a multivariate Gaussian distribution with zero mean and unit variance in `d` dimensions as a random walk proposal.

```julia
# Importing the sampling library
using AdvancedMH
rwmh = AdvancedMH.RWMH(d)
```

Sampling is then as easy as:

```julia
chain = sample(model, externalsampler(rwmh), 10_000)
```

## Going beyond the Turing API

As previously mentioned, the Turing wrappers can often limit the capabilities of the sampling libraries they wrap.
`AdvancedHMC`[^1] ((`AdvancedHMC`'s GitHub)[https://github.com/TuringLang/AdvancedHMC.jl]) is a clear example of this. A common practice when performing HMC is to provide an initial guess for the mass matrix.
However, the native HMC sampler within Turing only allows the user to specify the type of the mass matrix despite the two options being possible within `AdvancedHMC`.
Thankfully, we can use Turing's support for external samplers to define an HMC sampler with a custom mass matrix in `AdvancedHMC` and then use it to sample our Turing model.

We will use the library `Pathfinder`[^2] ((`Pathfinder`'s GitHub)[https://github.com/mlcolab/Pathfinder.jl]) to construct our estimate of mass matrix.
`Pathfinder` is a variational inference algorithm that first finds the maximum a posteriori (MAP) estimate of a target posterior distribution and then uses the trace of the optimization to construct a sequence of multivariate normal approximations to the target distribution.
In this process, `Pathfinder` computes an estimate of the mass matrix the user can access.

The code below shows this can be done in practice.

```julia
using AdvancedHMC, Pathfinder
# Running pathfinder
draws = 1_000
result_multi = multipathfinder(model, draws; nruns=8)

# Estimating the metric
inv_metric = result_multi.pathfinder_results[1].fit_distribution.Σ
metric = DenseEuclideanMetric(Matrix(inv_metric))

# Creating an AdvancedHMC NUTS sampler with the custom metric.
n_adapts = 1000 # Number of adaptation steps
tap = 0.9 # Large target acceptance probability to deal with the funnel structure of the posterior
nuts = AdvancedHMC.NUTS(tap; metric=metric)

# Sample
chain = sample(model, externalsampler(nuts), 10_000; n_adapts=1_000)
```

## Using new inference methods

So far we have used Turing's support for external samplers to go beyond the capabilities of the wrappers.
We want to use this support to employ a sampler not supported within Turing's ecosystem yet.
We will use the recently developed Micro-Cannoncial Hamiltonian Monte Carlo (MCHMC) sampler to showcase this.
MCHMC[^3,^4] ((MCHMC's GitHub)[https://github.com/JaimeRZP/MicroCanonicalHMC.jl]) is HMC sampler that uses one single Hamiltonian energy level to explore the whole parameter space.
This is achieved by simulating the dynamics of a microcanonical Hamiltonian with an additional noise term to ensure ergodicity.

Using this as well as other inference methods outside the Turing ecosystem is as simple as executing the code shown below:

```julia
using MicroCanonicalHMC
# Create MCHMC sampler
n_adapts = 1_000 # adaptation steps
tev = 0.01 # target energy variance
mchmc = MCHMC(n_adapts, tev; adaptive=true)

# Sample
chain = sample(model, externalsampler(mchmc), 10_000)
```

The only requirement to work with `externalsampler` is that the provided `sampler` must implement the AbstractMCMC.jl-interface [INSERT LINK] for a `model` of type `AbstractMCMC.LogDensityModel` [INSERT LINK].

As previously stated, in order to use external sampling libraries within `Turing` they must follow the `AbstractMCMC` API.
In this section, we will briefly dwell on what this entails.
First and foremost, the sampler should be a subtype of `AbstractMCMC.AbstractSampler`.
Second, the stepping function of the MCMC algorithm must be made defined using `AbstractMCMC.step` and follow the structure below:

```
# First step
function AbstractMCMC.step{T<:AbstractMCMC.AbstractSampler}(
rng::Random.AbstractRNG,
model::AbstractMCMC.LogDensityModel,
spl::T;
kwargs...,
)
[...]
return transition, sample
end

# N+1 step
function AbstractMCMC.step{T<:AbstractMCMC.AbstractSampler}(
rng::Random.AbstractRNG,
model::AbstractMCMC.LogDensityModel,
sampler::T,
state;
kwargs...,
)
[...]
return transition, sample
end
```

There are several characteristics to note in these functions:

- There must be two `step` functions:

+ A function that performs the first step and initializes the sampler.
+ A function that performs the following steps and takes an extra input, `state`, which carries the initialization information.

- The functions must follow the displayed signatures.
- The output of the functions must be a transition, the current state of the sampler, and a sample, what is saved to the MCMC chain.

The last requirement is that the transition must be structured with a field `θ` which contains the values of the parameters of the model for said transition.
This allows `Turing` to seamlessly extract the parameter values at each step of the chain when bundling the chains.
Note that if the external sampler produces transitions that Turing cannot parse the bundling of the samples will be different or fail.

For practical examples of how to adapt a sampling library to the `AbstractMCMC` interface, the readers can consult the following libraries:

- (AdvancedMH)[https://github.com/TuringLang/AdvancedMH.jl/blob/458a602ac32a8514a117d4c671396a9ba8acbdab/src/mh-core.jl#L73-L115]
- (AdvancedHMC)[https://github.com/TuringLang/AdvancedHMC.jl/blob/762e55f894d142495a41a6eba0eed9201da0a600/src/abstractmcmc.jl#L102-L170]
- (MicroCanonicalHMC)[https://github.com/JaimeRZP/MicroCanonicalHMC.jl/blob/master/src/abstractmcmc.jl] within `MicroCanonicalHMC`.

# Refences

[^1]: Xu et al, (AdvancedHMC.jl: A robust, modular and efficient implementation of advanced HMC algorithms)[http://proceedings.mlr.press/v118/xu20a/xu20a.pdf], 2019
[^2]: Zhang et al, (Pathfinder: Parallel quasi-Newton variational inference)[https://arxiv.org/abs/2108.03782], 2021
[^3]: Robnik et al, (Microcanonical Hamiltonian Monte Carlo)[https://arxiv.org/abs/2212.08549], 2022
[^4]: Robnik and Seljak, (Langevine Hamiltonian Monte Carlo)[https://arxiv.org/abs/2303.18221], 2023