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
There are a few tricky things to be aware of here. Firstly, if you are using Turing with FlexiChains, it's very likely that you are using FlexiChains v0.6. You should check this by inspecting your Project.toml or Manifest.toml: if this is indeed the case, please refer to the v0.6 documentation.
Secondly, although Pigeons has had a compatibility patch for DynamicPPL 0.41/0.42 merged, this has not yet been released. The current release of Pigeons.jl is not compatible with DynamicPPL v0.41 or v0.42, which are the versions currently used by Turing.jl. Although Pigeons has had a compatibility patch for DynamicPPL 0.41/0.42 merged, this has not yet been released.
Note that Pigeons' DynamicPPL extension requires DifferentiationInterface as one of its triggers. DI used to be a hard dependency of DynamicPPL, so this used to always be loaded, but since DynamicPPL v0.42 it is no longer a dependency. This means that you must explicitly load DI to use the Pigeons extension, even though we don't call any of its functions directly.
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, Distributionsimport DifferentiationInterfaceusing DynamicPPL: @model, @varname, filldist@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_successesendmy_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{AbstractPPL.VarName}.
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) ── AbstractPPL.VarName │
│ Base.ReshapedArray{Float64, 2, SubArray… p (1, 2) │
│ │
│ Extras (3) │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯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) ── AbstractPPL.VarName │
│ Matrix{Float64} p (1, 2) │
│ │
│ Extras (3) │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯Other models
Sampling with other models also works, and will return a VNChain (i.e., FlexiChain{VarNames.VarName}). Here is another example lifted from the Pigeons docs, where a custom log-density function is used.
using Randomstruct MyLogPotential n_trials::Int n_successes::Intendfunction (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)endPigeons.initialization(::MyLogPotential, ::Random.AbstractRNG, ::Int) = [0.5, 0.5]pt = pigeons(; target=MyLogPotential(100, 50), reference=MyLogPotential(0, 0), record=[traces],)PT(checkpoint = false, ...)By default, from_pigeons will return a VNChain with a single vector-valued parameter, @vn(param):
chn = FlexiChains.from_pigeons(pt)╭─FlexiChain (1024 iterations, 1 chain) ───────────────────────────────────────╮
│ ↓ iter = 1:1024 │
│ → chain = 1:1 │
│ │
│ Parameters (1) ── VarName │
│ 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
FlexiChain{AbstractPPL.VarName}. - Other models will produce
VNChain, i.e.,FlexiChain{VarNames.VarName}.
Other models
For non-DynamicPPL models, the parameter names are taken from Pigeons:
- If the target does not provide parameter names, Pigeons generates the names
param_1,param_2, etc. In this case, all the parameters are combined into a single vector-valued parameter,@vn(param). - Otherwise (e.g. for Stan models), each name that Pigeons returns becomes its own parameter. For example, a Stan model with parameters
p1andp2will give a chain with@vn(p1)and@vn(p2).
In both cases, the log-density is stored as Extra(:log_density).
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.