Skip to content

Commit

Permalink
Split up 'advanced usage' page
Browse files Browse the repository at this point in the history
  • Loading branch information
@penelopeysm
penelopeysm committed Sep 14, 2024
1 parent 0142603 commit cbb541d
Show file tree
Hide file tree
Showing 8 changed files with 284 additions and 247 deletions.
13 changes: 7 additions & 6 deletions _quarto.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,15 +59,15 @@ website:
- section: "Usage Tips"
collapse-level: 1
contents:
- text: "Mode Estimation"
href: tutorials/docs-17-mode-estimation/index.qmd
- tutorials/docs-09-using-turing-advanced/index.qmd
- tutorials/docs-10-using-turing-autodiff/index.qmd
- tutorials/usage-custom-distribution/index.qmd
- tutorials/usage-modifying-logprob/index.qmd
- tutorials/usage-generated-quantities/index.qmd
- tutorials/docs-17-mode-estimation/index.qmd
- tutorials/docs-13-using-turing-performance-tips/index.qmd
- tutorials/docs-11-using-turing-dynamichmc/index.qmd
- tutorials/docs-15-using-turing-sampler-viz/index.qmd
- text: "External Samplers"
href: tutorials/docs-16-using-turing-external-samplers/index.qmd
- tutorials/docs-11-using-turing-dynamichmc/index.qmd
- tutorials/docs-16-using-turing-external-samplers/index.qmd

- section: "Tutorials"
contents:
Expand Down Expand Up @@ -107,6 +107,7 @@ website:
- section: "DynamicPPL in Depth"
collapse-level: 1
contents:
- tutorials/dev-model-manual/index.qmd
- tutorials/docs-05-for-developers-compiler/index.qmd
- text: "A Mini Turing Implementation I: Compiler"
href: tutorials/14-minituring/index.qmd
Expand Down
59 changes: 59 additions & 0 deletions tutorials/dev-model-manual/index.qmd
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
---
title: Manually Defining a Model
engine: julia
---

Traditionally, models in Turing are defined using the `@model` macro:

```{julia}
using Turing
@model function gdemo(x)
# Set priors.
s² ~ InverseGamma(2, 3)
m ~ Normal(0, sqrt(s²))
# Observe each value of x.
@. x ~ Normal(m, sqrt(s²))
end
model = gdemo([1.5, 2.0])
```

The `@model` macro accepts a function definition and rewrites it such that call of the function generates a `Model` struct for use by the sampler.

However, models can be constructed by hand without the use of a macro.
Taking the `gdemo` model above as an example, the macro-based definition can be implemented also (a bit less generally) with the macro-free version

```{julia}
# Create the model function.
function gdemo2(model, varinfo, context, x)
# Assume s² has an InverseGamma distribution.
s², varinfo = DynamicPPL.tilde_assume!!(
context, InverseGamma(2, 3), Turing.@varname(s²), varinfo
)
# Assume m has a Normal distribution.
m, varinfo = DynamicPPL.tilde_assume!!(
context, Normal(0, sqrt(s²)), Turing.@varname(m), varinfo
)
# Observe each value of x[i] according to a Normal distribution.
return DynamicPPL.dot_tilde_observe!!(
context, Normal(m, sqrt(s²)), x, Turing.@varname(x), varinfo
)
end
gdemo2(x) = Turing.Model(gdemo2, (; x))
# Instantiate a Model object with our data variables.
model2 = gdemo2([1.5, 2.0])
```

We can sample from this model in the same way:

```{julia}
setprogress!(false)
chain = sample(model2, NUTS(), 1000)
```

The subsequent pages in this section will show how the `@model` macro does this behind-the-scenes.
243 changes: 5 additions & 238 deletions tutorials/docs-09-using-turing-advanced/index.qmd
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,242 +3,9 @@ title: Advanced Usage
engine: julia
---

```{julia}
#| echo: false
#| output: false
using Pkg;
Pkg.instantiate();
```
This page has been separated into new sections. Please update any bookmarks you might have:

```{julia}
#| echo: false
using Distributions, Turing, Random, Bijectors
```

## How to Define a Customized Distribution

`Turing.jl` supports the use of distributions from the Distributions.jl package. By extension, it also supports the use of customized distributions by defining them as subtypes of `Distribution` type of the Distributions.jl package, as well as corresponding functions.

Below shows a workflow of how to define a customized distribution, using our own implementation of a simple `Uniform` distribution as a simple example.

### 1. Define the Distribution Type

First, define a type of the distribution, as a subtype of a corresponding distribution type in the Distributions.jl package.

```{julia}
struct CustomUniform <: ContinuousUnivariateDistribution end
```

### 2. Implement Sampling and Evaluation of the log-pdf

Second, define `rand` and `logpdf`, which will be used to run the model.

```{julia}
# sample in [0, 1]
Distributions.rand(rng::AbstractRNG, d::CustomUniform) = rand(rng)
# p(x) = 1 → logp(x) = 0
Distributions.logpdf(d::CustomUniform, x::Real) = zero(x)
```

### 3. Define Helper Functions

In most cases, it may be required to define some helper functions.

#### 3.1 Domain Transformation

Certain samplers, such as `HMC`, require the domain of the priors to be unbounded. Therefore, to use our `CustomUniform` as a prior in a model we also need to define how to transform samples from `[0, 1]` to `ℝ`. To do this, we simply need to define the corresponding `Bijector` from `Bijectors.jl`, which is what `Turing.jl` uses internally to deal with constrained distributions.

To transform from `[0, 1]` to `ℝ` we can use the `Logit` bijector:

```{julia}
Bijectors.bijector(d::CustomUniform) = Logit(0.0, 1.0)
```

You'd do the exact same thing for `ContinuousMultivariateDistribution` and `ContinuousMatrixDistribution`. For example, `Wishart` defines a distribution over positive-definite matrices and so `bijector` returns a `PDBijector` when called with a `Wishart` distribution as an argument. For discrete distributions, there is no need to define a bijector; the `Identity` bijector is used by default.

Alternatively, for `UnivariateDistribution` we can define the `minimum` and `maximum` of the distribution

```{julia}
Distributions.minimum(d::CustomUniform) = 0.0
Distributions.maximum(d::CustomUniform) = 1.0
```

and `Bijectors.jl` will return a default `Bijector` called `TruncatedBijector` which makes use of `minimum` and `maximum` derive the correct transformation.

Internally, Turing basically does the following when it needs to convert a constrained distribution to an unconstrained distribution, e.g. when sampling using `HMC`:

```{julia}
dist = Gamma(2,3)
b = bijector(dist)
transformed_dist = transformed(dist, b) # results in distribution with transformed support + correction for logpdf
```

and then we can call `rand` and `logpdf` as usual, where

- `rand(transformed_dist)` returns a sample in the unconstrained space, and
- `logpdf(transformed_dist, y)` returns the log density of the original distribution, but with `y` living in the unconstrained space.

To read more about Bijectors.jl, check out [the project README](https://github.com/TuringLang/Bijectors.jl).

## Update the accumulated log probability in the model definition

Turing accumulates log probabilities internally in an internal data structure that is accessible through
the internal variable `__varinfo__` inside of the model definition (see below for more details about model internals).
However, since users should not have to deal with internal data structures, a macro `Turing.@addlogprob!` is provided
that increases the accumulated log probability. For instance, this allows you to
[include arbitrary terms in the likelihood](https://github.com/TuringLang/Turing.jl/issues/1332)

```{julia}
using Turing
myloglikelihood(x, μ) = loglikelihood(Normal(μ, 1), x)
@model function demo(x)
μ ~ Normal()
Turing.@addlogprob! myloglikelihood(x, μ)
end
```

and to [reject samples](https://github.com/TuringLang/Turing.jl/issues/1328):

```{julia}
using Turing
using LinearAlgebra
@model function demo(x)
m ~ MvNormal(zero(x), I)
if dot(m, x) < 0
Turing.@addlogprob! -Inf
# Exit the model evaluation early
return nothing
end
x ~ MvNormal(m, I)
return nothing
end
```

Note that `@addlogprob!` always increases the accumulated log probability, regardless of the provided
sampling context. For instance, if you do not want to apply `Turing.@addlogprob!` when evaluating the
prior of your model but only when computing the log likelihood and the log joint probability, then you
should [check the type of the internal variable `__context_`](https://github.com/TuringLang/DynamicPPL.jl/issues/154)
such as

```{julia}
#| eval: false
if DynamicPPL.leafcontext(__context__) !== Turing.PriorContext()
Turing.@addlogprob! myloglikelihood(x, μ)
end
```

## Model Internals

The `@model` macro accepts a function definition and rewrites it such that call of the function generates a `Model` struct for use by the sampler.
Models can be constructed by hand without the use of a macro.
Taking the `gdemo` model as an example, the macro-based definition

```{julia}
using Turing
@model function gdemo(x)
# Set priors.
s² ~ InverseGamma(2, 3)
m ~ Normal(0, sqrt(s²))
# Observe each value of x.
@. x ~ Normal(m, sqrt(s²))
end
model = gdemo([1.5, 2.0])
```

can be implemented also (a bit less generally) with the macro-free version

```{julia}
using Turing
# Create the model function.
function gdemo(model, varinfo, context, x)
# Assume s² has an InverseGamma distribution.
s², varinfo = DynamicPPL.tilde_assume!!(
context, InverseGamma(2, 3), Turing.@varname(s²), varinfo
)
# Assume m has a Normal distribution.
m, varinfo = DynamicPPL.tilde_assume!!(
context, Normal(0, sqrt(s²)), Turing.@varname(m), varinfo
)
# Observe each value of x[i] according to a Normal distribution.
return DynamicPPL.dot_tilde_observe!!(
context, Normal(m, sqrt(s²)), x, Turing.@varname(x), varinfo
)
end
gdemo(x) = Turing.Model(gdemo, (; x))
# Instantiate a Model object with our data variables.
model = gdemo([1.5, 2.0])
```

### Reparametrization and generated_quantities

Often, the most natural parameterization for a model is not the most computationally feasible. Consider the following
(efficiently reparametrized) implementation of Neal's funnel [(Neal, 2003)](https://arxiv.org/abs/physics/0009028):

```{julia}
#| eval: false
@model function Neal()
# Raw draws
y_raw ~ Normal(0, 1)
x_raw ~ arraydist([Normal(0, 1) for i in 1:9])
# Transform:
y = 3 * y_raw
x = exp.(y ./ 2) .* x_raw
# Return:
return [x; y]
end
```

In this case, the random variables exposed in the chain (`x_raw`, `y_raw`) are not in a helpful form — what we're after is the deterministically transformed variables `x, y`.

More generally, there are often quantities in our models that we might be interested in viewing, but which are not explicitly present in our chain.

We can generate draws from these variables — in this case, `x, y` — by adding them as a return statement to the model, and then calling `generated_quantities(model, chain)`. Calling this function outputs an array of values specified in the return statement of the model.

For example, in the above reparametrization, we sample from our model:

```{julia}
#| eval: false
chain = sample(Neal(), NUTS(), 1000)
```

and then call:

```{julia}
#| eval: false
generated_quantities(Neal(), chain)
```

to return an array for each posterior sample containing `x1, x2, ... x9, y`.

In this case, it might be useful to reorganize our output into a matrix for plotting:

```{julia}
#| eval: false
reparam_chain = reduce(hcat, generated_quantities(Neal(), chain))'
```

Where we can recover a vector of our samples as follows:

```{julia}
#| eval: false
x1_samples = reparam_chain[:, 1]
y_samples = reparam_chain[:, 10]
```

## Task Copying

Turing [copies](https://github.com/JuliaLang/julia/issues/4085) Julia tasks to deliver efficient inference algorithms, but it also provides alternative slower implementation as a fallback. Task copying is enabled by default. Task copying requires us to use the `TapedTask` facility which is provided by [Libtask](https://github.com/TuringLang/Libtask.jl) to create tasks.
- [Custom Distributions](../../tutorials/usage-custom-distribution)
- [Modifying the Log Probability](../../tutorials/usage-modifying-logprob/)
- [Defining a Model without `@model`](../../tutorials/dev-model-manual/)
- [Reparametrization and Generated Quantities](../../tutorials/usage-generated-quantities/)
2 changes: 1 addition & 1 deletion tutorials/docs-15-using-turing-sampler-viz/index.qmd
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,4 +194,4 @@ Next, we plot using 50 particles.
```{julia}
c = sample(model, PG(50), 1000)
plot_sampler(c)
```
```
4 changes: 2 additions & 2 deletions tutorials/docs-16-using-turing-external-samplers/index.qmd
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
---
title: Using External Sampler
title: Using External Samplers
engine: julia
---

Expand Down Expand Up @@ -176,4 +176,4 @@ For practical examples of how to adapt a sampling library to the `AbstractMCMC`
[^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
[^4]: Robnik and Seljak, [Langevine Hamiltonian Monte Carlo](https://arxiv.org/abs/2303.18221), 2023
Loading

0 comments on commit cbb541d

Please sign in to comment.