Why FlexiChains?
If you are a Turing user, you may well be asking what benefits FlexiChains.jl offers over the previous default of MCMCChains.jl.
In one sentence: FlexiChains is a more faithful representation of data. MCMCChains.jl places extremely strong restrictions on its data structure, which leads to an irrevocable loss of information.
Of course, on its own this doesn't mean much (unless you are a software engineering purist!). So this page will demonstrate a few concrete scenarios where FlexiChains has a practical advantage.
Heterogeneous parameter types
using Random, Turing, FlexiChains, MCMCChains
@model function f()
x ~ Normal()
y ~ Poisson(3.0)
endf (generic function with 2 methods)When sampling from this model, one should expect that the samples of x are stored as floats, whereas the samples of y are stored as integers, because that is what these distributions produce.
Under the hood, MCMCChains stores the values of all parameters in a single array, which means that all samples get converted into the same type.
mchain = sample(Xoshiro(468), f(), MH(), 50; chain_type=MCMCChains.Chains);
fchain = sample(Xoshiro(468), f(), MH(), 50; chain_type=FlexiChains.VNChain);
(eltype(mchain[:y]), eltype(fchain[:y]))(Float64, Int64)Now, Distributions.jl kindly allows you to call logpdf(::Poisson, x::Float64) and returns the correct value if isinteger(x). So, if your model is simple enough, you can still use functions such as returned and predict without running into errors, even with MCMCChains.
However, if you attempt to use the value of x somewhere where it needs to be an integer, such as...
@model function f2()
# A more realistic scenario is `n ~ Poisson(...)`.
# That errors with MCMCChains for a *different* reason;
# we'll come to that in a while.
n ~ DiscreteUniform(2, 2)
x ~ MvNormal(zeros(n), I)
return sum(x)
end
mchain = sample(Xoshiro(468), f2(), MH(), 50; chain_type=MCMCChains.Chains);Chains MCMC chain (50×7×1 Array{Float64, 3}):
Iterations = 1:1:50
Number of chains = 1
Samples per chain = 50
Wall duration = 0.28 seconds
Compute duration = 0.28 seconds
parameters = n, x[1], x[2]
internals = accepted, logprior, loglikelihood, logjoint
Use `describe(chains)` for summary statistics and quantiles.you will find that it errors, because n is stored as 2.0 in the chain:
#! format: off
returned(f2(), mchain)MethodError: no method matching zeros(::Float64)
The function `zeros` exists, but no method is defined for this combination of argument types.
Closest candidates are:
zeros(::Type{SA}) where {T, SA<:(StaticArraysCore.StaticArray{<:Tuple, T})}
@ StaticArrays ~/.julia/packages/StaticArrays/EnGcj/src/arraymath.jl:2
zeros(::Type{SA}) where SA<:(StaticArraysCore.StaticArray)
@ StaticArrays ~/.julia/packages/StaticArrays/EnGcj/src/arraymath.jl:1
zeros(::Type{MonteCarloMeasurements.StaticParticles{T, N}}, ::Integer) where {T, N}
@ MonteCarloMeasurements ~/.julia/packages/MonteCarloMeasurements/zXluZ/src/particles.jl:343
...Now, you could work around this with zeros(Int(n)), but that's deeply unsatisfying, because n really should be an integer. Good news: FlexiChains will store it as an integer for you!
fchain = sample(Xoshiro(468), f2(), MH(), 50; chain_type=FlexiChains.VNChain)
returned(f2(), fchain)┌ 50×1 DimArray{Float64, 2} ┐
├───────────────────────────┴────────────────────────── dims ┐
↓ iter Sampled{Int64} 1:50 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:1 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────────┘
↓ → 1
1 0.558732
2 1.40841
3 1.36466
4 3.21083
⋮
47 1.16891
48 -0.067146
49 0.199018
50 -0.921291Missing data
As a variant on the above bug with proper type representation, consider the case where some data is missing:
using Random, Turing, FlexiChains, MCMCChains
@model function f()
x ~ Normal()
if x > 0
y ~ Normal(x)
end
end
mchain = sample(Xoshiro(468), f(), MH(), 50; chain_type=MCMCChains.Chains)
fchain = sample(Xoshiro(468), f(), MH(), 50; chain_type=FlexiChains.VNChain)╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 1:50 │
│ → chain = 1:1 │
│ │
│ Parameters (2) ── VarName │
│ Float64 x │
│ Union{Missing, Float64} y │
│ │
│ Extras (4) │
│ Bool accepted │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯In some samples, y will be present and not in others. Because MCMCChains forces all parameters to be in the same array, this means that the entire array must have an element type of Union{Missing, Float64}. With MCMCChains this gets propagated to all parameters, even those that are never missing, such as x.
(eltype(mchain[:x]), eltype(fchain[:x]))(Union{Missing, Float64}, Float64)Okay, maybe you don't ever have such weird models. It turns out though that you can still run into this. In Turing's MCMC sampling, the first step is not an actual MCMC step, but rather just the initial parameters (either sampled or provided by the user). Thus, there are no 'sampler statistics' for the first step, and these are stored as missing in MCMCChains. That means that all the parameters become Union{Missing, Float64}!
Variable-length parameters
Above we used an example with DiscreteUniform(2, 2) to demonstrate what happens when Int-valued parameters get converted to Floats. This is of course a bit pointless since sampling from that always gives 2. Let's see now what happens when we have a truly variable-length parameter:
using Random, Turing, FlexiChains, MCMCChains
@model function varlen()
n ~ Poisson(3.5)
x ~ MvNormal(zeros(n), I)
y ~ Normal(sum(x))
end
model = varlen()
cond_model = varlen() | (; y=2.0)
mchain = sample(Xoshiro(468), cond_model, MH(), 50; chain_type=MCMCChains.Chains)
fchain = sample(Xoshiro(468), cond_model, MH(), 50; chain_type=FlexiChains.VNChain)╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 1:50 │
│ → chain = 1:1 │
│ │
│ Parameters (2) ── VarName │
│ Int64 n │
│ Vector{Float64} x │
│ │
│ Extras (4) │
│ Bool accepted │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯So far, so good; we can sample from everything just fine. The trouble comes when you want to use something like predict or returned which involves feeding the samples from the chain back into the model.
#! format: off
returned(f2(), mchain)UndefVarError: `f2` not defined in `Main`
Suggestion: check for spelling errors or missing imports.Warning
The above will probably fail, but it can actually run successfully, if you are lucky enough to get a sample where n is larger than or equal to that in all the other samples. If the built docs don't show an error, try running it in the REPL, and you'll find that it will almost always fail.
The reason why MCMCChains fails here is because it does two things: 2. It stores x as a series of elements x[1], x[2], and so on.
- When reconstructing the value of
xfor use in the model, it doesn't know how longxis supposed to be. It determines this by running the model once, and taking the value ofxfrom that specific run of the model.
This of course ignores the fact that x can have different lengths.
In contrast, FlexiChains does two things: 2. Where possible, it will store x as a single parameter.
- As a guard against situations where this isn't possible (e.g. if not all values of
xare filled in), it also stores the structure ofxas part of the chain, so that it can always reconstruct it correctly.
This means that regardless of what value n takes in the samples, predict and returned will always work correctly with FlexiChains.
predict(model, fchain)╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 1:50 │
│ → chain = 1:1 │
│ │
│ Parameters (3) ── VarName │
│ Int64 n │
│ Vector{Float64} x │
│ Float64 y │
│ │
│ Extras (4) │
│ Bool accepted │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯Arbitrary types
Turing provides this very nice operator :=, which lets you store arbitrary values in the chain during an MCMC run.
The problem with MCMCChains is, as ever, you can only store things that are Real or some array thereof. Even a simple string will fail:
using Random, Turing, FlexiChains, MCMCChains
@model function hasstring()
x ~ Normal()
y := "$x"
end
#! format: off
mchain = sample(Xoshiro(468), hasstring(), MH(), 50; chain_type=MCMCChains.Chains)[ Info: [Turing]: progress logging is disabled globally
MethodError: no method matching MCMCChains.Chains(::Array{Any, 3}, ::Vector{Symbol}, ::@NamedTuple{internals::Vector{Symbol}}; info::@NamedTuple{varname_to_symbol::OrderedCollections.OrderedDict{VarName, Symbol}})
The type `MCMCChains.Chains` exists, but no method is defined for this combination of argument types when trying to construct it.
Closest candidates are:
MCMCChains.Chains(::AbstractArray{<:Union{Missing, Real}, 3}, ::AbstractVector{Symbol}, ::Any; start, thin, iterations, evidence, info)
@ MCMCChains ~/.julia/packages/MCMCChains/vqJuv/src/chains.jl:28
MCMCChains.Chains(::AbstractVector{<:AbstractVector{<:Union{Missing, Real}}}, ::Any...; kwargs...)
@ MCMCChains ~/.julia/packages/MCMCChains/vqJuv/src/chains.jl:6
MCMCChains.Chains(::AbstractVector{<:Union{Missing, Real}}, ::Any...; kwargs...)
@ MCMCChains ~/.julia/packages/MCMCChains/vqJuv/src/chains.jl:10
...FlexiChains will let you store anything you like! This includes strings, but also more complex objects such as the output of a differential equation solver, or customised structs, as demonstrated in this Gaussian process tutorial.
fchain = sample(Xoshiro(468), hasstring(), MH(), 50; chain_type=FlexiChains.VNChain)╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 1:50 │
│ → chain = 1:1 │
│ │
│ Parameters (2) ── VarName │
│ Float64 x │
│ String y │
│ │
│ Extras (4) │
│ Bool accepted │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯Reconstructing parameters
Suppose you have some array-valued parameter.
using Random, Turing, FlexiChains, MCMCChains
@model lkj() = x ~ LKJCholesky(3, 2.0)
mchain = sample(Xoshiro(468), lkj(), NUTS(), 50; chain_type=MCMCChains.Chains);
fchain = sample(Xoshiro(468), lkj(), NUTS(), 50; chain_type=FlexiChains.VNChain);╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 26:75 │
│ → chain = 1:1 │
│ │
│ Parameters (1) ── VarName │
│ LinearAlgebra.Cholesky{Float64, Matrix{… x │
│ │
│ Extras (14) │
│ Int64 n_steps, tree_depth │
│ Bool is_accept, numerical_error │
│ Float64 acceptance_rate, log_density, hamiltonian_energy, │
│ hamiltonian_energy_error, max_hamiltonian_energy_error, step_size, │
│ nom_step_size, logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯With FlexiChains the Cholesky samples are kept together:
fchain[@varname(x)][iter=1, chain=1]LinearAlgebra.Cholesky{Float64, Matrix{Float64}}
L factor:
3×3 LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}:
1.0 ⋅ ⋅
-0.104367 0.994539 ⋅
0.122204 -0.599314 0.791131Because MCMCChains stores all its data in a single array, it has to flatten this parameter, so good luck trying to reconstruct it.
(mchain[Symbol("x.L[1, 1]")][1, 1], mchain[Symbol("x.L[2, 1]")][1, 1]) # ...(1.0, -0.10436666574417859)VarNames as keys
Did you notice in that last line we had to write something like Symbol("x.L[1, 1]")? And since it's a Symbol, we have to get the name exactly right, we couldn't do (for example) x.L[1,1] without the space after the comma?
That's because MCMCChains uses AxisArrays.jl under the hood, which allows you to index into the chain using Symbols — but only Symbols. FlexiChains retains the original VarNames used by Turing, which is a far richer type and allows you to use keys that actually carry meaning, rather than just being strings that have to match exactly.
fchain[@varname(x.L[1, 1])][iter=1, chain=1] # No space!1.0fchain[@varname(x.L[:, 1])][iter=1, chain=1] # Index into `x` any way you like3-element Vector{Float64}:
1.0
-0.10436666574417859
0.12220440942739924If you have x stored as a full vector, you can also do fancy things like indexing into x[end], which will give you the last element of x regardless of how long it is in that particular sample.
@model function varlen_again()
n ~ Poisson(3.5)
x ~ MvNormal(zeros(n), I)
end
fchain = sample(Xoshiro(468), varlen_again(), MH(), 5; chain_type=FlexiChains.VNChain);
fchain[@varname(x)]┌ 5×1 DimArray{Vector{Float64}, 2} Parameter(x) ┐
├───────────────────────────────────────────────┴───── dims ┐
↓ iter Sampled{Int64} 1:5 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:1 ForwardOrdered Regular Points
└───────────────────────────────────────────────────────────┘
↓ → 1
1 [-0.203115, 0.119206, -0.371982, 1.50102]
2 [-0.298648, 2.15728, -0.810716, 0.648159]
3 [0.171813, -1.07679, -1.50273, -1.05637]
4 [-1.06928, 0.435984, -1.00153, -0.74393]
5 [2.05543, -0.555033, -1.07123]fchain[@varname(x[end])]┌ 5×1 DimArray{Float64, 2} Parameter(x[DynamicIndex(end)]) ┐
├──────────────────────────────────────────────────────────┴ dims ┐
↓ iter Sampled{Int64} 1:5 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:1 ForwardOrdered Regular Points
└─────────────────────────────────────────────────────────────────┘
↓ → 1
1 1.50102
2 0.648159
3 -1.05637
4 -0.74393
5 -1.07123Interoperability with the rest of Turing
Suppose you have sampled a chain, and you want to use something from it as the starting point for a new chain, or a new optimisation, or something.
using Random, Turing, FlexiChains, MCMCChains
@model function twonorm()
x ~ Normal()
y ~ Normal()
return x + y
end
mchain = sample(Xoshiro(468), twonorm(), MH(), 50; chain_type=MCMCChains.Chains);
fchain = sample(Xoshiro(468), twonorm(), MH(), 50; chain_type=FlexiChains.VNChain);╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 1:50 │
│ → chain = 1:1 │
│ │
│ Parameters (2) ── VarName │
│ Float64 x, y │
│ │
│ Extras (4) │
│ Bool accepted │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯Let's say you want to use the last sample of x as the starting point for another sampler.
Array(mchain)[50, :, :]2×1 Matrix{Float64}:
-1.5640135830760107
-1.4689274565945687Even at this early point, it is already quite ugly: you have to know that the samples are stored in a 3D array, and that the first dimension corresponds to iterations, the second to parameters, and the third to chains. Secondly, it gives you a vector of parameters, and it's nontrivial to figure out from this how these map to the original variables in the model. You can use
names(MCMCChains.get_sections(mchain, :parameters))2-element Vector{Symbol}:
:x
:yand for this trivial model it's very clear that the first one is x and the second is y. However, this generalises to complex samples very poorly (consider e.g. matrices and Cholesky samples again), and is very difficult to write in a programmatic way!
In contrast, with FlexiChains, you can just do
vnt = FlexiChains.parameters_at(fchain, 50, 1)VarNamedTuple
├─ x => -1.5640135830760107
└─ y => -1.4689274565945687which directly gives you a VarNamedTuple that will 'just work' with the rest of Turing. For example, you can use it as the starting point for a new sampler:
init = InitFromParams(vnt)
sample(twonorm(), NUTS(), 50; initial_params=init, chain_type=FlexiChains.VNChain)╭─FlexiChain (50 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter = 26:75 │
│ → chain = 1:1 │
│ │
│ Parameters (2) ── VarName │
│ Float64 x, y │
│ │
│ Extras (14) │
│ Int64 n_steps, tree_depth │
│ Bool is_accept, numerical_error │
│ Float64 acceptance_rate, log_density, hamiltonian_energy, │
│ hamiltonian_energy_error, max_hamiltonian_energy_error, step_size, │
│ nom_step_size, logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯or use it to begin an optimisation (a bit pointless for our trivial model, but you get the idea!):
maximum_a_posteriori(twonorm(); initial_params=init)ModeResult
├ estimator : Turing.Optimisation.MAP
├ lp : -1.8378770664093456
├ params : VarNamedTuple with 2 entries
│ ├ x => 0.0
│ └ y => 0.0
│ linked : true
└ (2 more fields: optim_result, ldf)or pass it to a function like returned:
returned(twonorm(), vnt)-3.0329410396705794Performance (on important things)
In the following model, y is a single parameter that is a vector of length N. That means that when you use functions like returned or predict on a chain, MCMCChains has to somehow reconstruct the vector y from its components which are all stored separately.
using Turing, FlexiChains, MCMCChains, Random
@model function longvec(N)
m ~ Normal(0)
y ~ filldist(Normal(m), N)
endlongvec (generic function with 2 methods)It turns out that if N is small, MCMCChains does just fine, and is in fact even faster than FlexiChains (because constructing a FlexiChain has more overhead).
using Chairmarks: @be
function benchmark(N)
model = (longvec(N) | (y=rand(Xoshiro(468), Normal(2.0), N),))
mchain = sample(Xoshiro(468), model, NUTS(), 500; chain_type=MCMCChains.Chains)
fchain = sample(Xoshiro(468), model, NUTS(), 500; chain_type=FlexiChains.VNChain)
mt = @be predict(longvec(N), mchain)
ft = @be predict(longvec(N), fchain)
return (N=N, mcmcchains=median(mt).time, flexichains=median(ft).time)
end
benchmark(2) # The results are in seconds(N = 2, mcmcchains = 0.0035007740000000003, flexichains = 0.003316152)But MCMCChains scales really poorly with N.
benchmark(1000)(N = 1000, mcmcchains = 0.35765013, flexichains = 0.292193497)Avoiding name clashes
All of Turing.jl's samplers include some 'sampler statistics' in the output chain. These are pretty useful things like the step size, whether a transition was accepted, the log-probabilities, and so on.
But if you have a parameter that just happens to share a name with these, then MCMCChains will make it pretty hard for you to get one of them.
using Random, Turing, FlexiChains, MCMCChains
# Oops! This will clash with the actual log prior.
@model pr() = logprior ~ Normal()
mchain = sample(Xoshiro(468), pr(), MH(), 50; chain_type=MCMCChains.Chains);
collect(keys(mchain))5-element Vector{Symbol}:
:logprior
:accepted
:logprior
:loglikelihood
:logjointThere are two columns labelled :logprior. Of course, one is your parameter, the other is the actual log prior probability. It's a mystery which one is which, and which you get when you do mchain[:logprior]! You could avoid this if you knew exactly which keys the sampler returns, but in general this isn't documented anywhere. (It should be, of course.)
FlexiChains avoids clashes by completely separating Parameter and Extra keys, meaning that you can use any name you like without worrying about a rogue sampler breaking your workflow.
DimensionalData.jl indexing
As you will have noticed, FlexiChains uses DimensionalData.jl to return information-rich matrices. That means that you can use all the selectors from DimensionalData.jl to extract exactly what you want. Don't want to use those? No problem; good old 1-based indices work fine too.