Pigeons.jl

Documentation for Pigeons.jl ↗

Pigeons.jl is a package that provides algorithms such as parallel tempering for challenging posterior distributions. FlexiChains provides a function, FlexiChains.from_pigeons, to convert a Pigeons.PT object into a FlexiChain. This works with any model that Pigeons supports, including DynamicPPL models and others.

Turing.jl models

Compatibility

The current version of FlexiChains only works with DynamicPPL 0.41, whereas the current release of Pigeons only works with DynamicPPL 0.40. To work around this, you can check out the code from this PR, which allows you to use DynamicPPL 0.41 (but not 0.42 yet).

The following example uses a Turing model as the log-density function (the exact model is modified slightly from the Pigeons docs):

using Pigeons, FlexiChains, Turing

@model function my_turing_model(n_trials, n_successes)
    p ~ filldist(Uniform(0, 1), 1, 2)
    n_successes ~ Binomial(n_trials, prod(p))
    return n_successes
end

my_turing_model (generic function with 2 methods)

When sampling, you have to specify record=[traces] so that the actual samples are stored in the returned Pigeons.PT struct. This object can then be converted to a FlexiChain. Specifically, with an underlying Turing model, a VNChain will be returned. This has all the usual benefits of a FlexiChain: for example, p is stored as a 1 × 2 matrix.

my_turing_target = TuringLogPotential(my_turing_model(100, 50))
pt = pigeons(; target=my_turing_target, record=[traces])
chn = FlexiChains.from_pigeons(pt)

╭─FlexiChain (1024 iterations, 1 chain) ───────────────────────────────────────╮
│ ↓ iter  = 1:1024                                                             │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── VarName                                                    │
│  Base.ReshapedArray{Float64, 2, SubArray…  p (1, 2)                          │
│                                                                              │
│ Extras (3)                                                                   │
│  Float64  logprior, loglikelihood, logjoint                                  │
╰──────────────────────────────────────────────────────────────────────────────╯

Note

Technically, it's stored as a lazy reshape/view of a 1 × 2 matrix. This is because of performance optimisations in Bijectors.jl and DynamicPPL, which avoid materialising the actual matrix unless really needed. If for any reason you want to materialise it, you can use FlexiChains.transform_values to apply collect to each sample:

chn2 = FlexiChains.transform_values(chn, @varname(p) => collect)

╭─FlexiChain (1024 iterations, 1 chain) ───────────────────────────────────────╮
│ ↓ iter  = 1:1024                                                             │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── VarName                                                    │
│  Matrix{Float64}  p (1, 2)                                                   │
│                                                                              │
│ Extras (3)                                                                   │
│  Float64  logprior, loglikelihood, logjoint                                  │
╰──────────────────────────────────────────────────────────────────────────────╯

Other models

Sampling with other models also works, but will return a SymChain (i.e., FlexiChain{Symbol}) instead of VNChain. Here is another example lifted from the Pigeons docs, where a custom log-density function is used:

using Random

struct MyLogPotential
    n_trials::Int
    n_successes::Int
end
function (log_potential::MyLogPotential)(x)
    p1, p2 = x
    ((0 < p1 < 1) && (0 < p2 < 1)) || return -Inf
    logpdf(Binomial(log_potential.n_trials, p1 * p2), log_potential.n_successes)
end
Pigeons.initialization(::MyLogPotential, ::Random.AbstractRNG, ::Int) = [0.5, 0.5]
pt = pigeons(;
    target=MyLogPotential(100, 50),
    reference=MyLogPotential(0, 0),
    record=[traces],
)

chn = FlexiChains.from_pigeons(pt)

╭─FlexiChain (1024 iterations, 1 chain) ───────────────────────────────────────╮
│ ↓ iter  = 1:1024                                                             │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── Symbol                                                     │
│  Vector{Float64}  param (2,)                                                 │
│                                                                              │
│ Extras (1)                                                                   │
│  Float64  log_density                                                        │
╰──────────────────────────────────────────────────────────────────────────────╯

Docstrings

FlexiChains.from_pigeons Function

FlexiChains.from_pigeons(pt::Pigeons.PT)

Convert the result of a Pigeons.jl sampling run into a FlexiChain.

The run must have been performed with pigeons(; ... record = [traces, ...]) so that the samples can be obtained. If this was not included, then this function will error.

The output type depends on the kind of model sampled from:

• DynamicPPL models will produce VNChain, i.e., FlexiChain{VarName}.

• Other models will produce SymChain, i.e., FlexiChain{Symbol}

DynamicPPL models

Note that for DynamicPPL models, calling from_pigeons will result in the model being reevaluated in order to recover the structure of the parameters. This is necessary because Pigeons stores a flattened version of the parameters. It should generally be the case that the time taken for this reevaluation is negligible compared to the actual sampling time.

The reevaluation isn't without benefit, though: you will also obtain (for free) the breakdown of the log-density between prior and likelihood, plus any x := expr variables in the model, which Pigeons does not save.

source