AbstractMCMC.jl samplers

If you have written a sampler that uses the AbstractMCMC interface, then you can get 'free' support for bundling into a FlexiChain{VarName} and a FlexiChain{Symbol} respectively by overloading FlexiChains.to_vnt_and_stats and FlexiChains.to_nt_and_stats for your sampler's transition type.

Specifically, your sampler must implement AbstractMCMC.step, which returns a tuple of (transition, state). The first of these two objects is what is passed to the functions above.

As a concrete example, let's use the following setup:

using FlexiChains: FlexiChains, FlexiChain
using AbstractMCMC

struct M <: AbstractMCMC.AbstractModel end
struct S <: AbstractMCMC.AbstractSampler end
struct T end
function AbstractMCMC.step(rng, ::M, ::S, state=nothing; kwargs...)
    T(), nothing
end

Here T represents what AbstractMCMC calls a 'transition': it is the first return value of AbstractMCMC.step. Of course, in practice, your transition will carry more information than this.

Now suppose you want to sample with this and obtain a FlexiChain{Symbol}. You should then overload to_nt_and_stats to return a tuple of two NamedTuples, the first being the parameter values, and the second being any Extras to include in the chain:

FlexiChains.to_nt_and_stats(::T) = ((; hello=1.0), (; world=2.0))

Then you can do

sample(M(), S(), 10; chain_type=FlexiChain{Symbol})

╭─FlexiChain (10 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter  = 1:10                                                               │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── Symbol                                                     │
│  Float64  hello                                                              │
│                                                                              │
│ Extras (1)                                                                   │
│  Float64  world                                                              │
╰──────────────────────────────────────────────────────────────────────────────╯

Likewise, you can overload to_vnt_and_stats to obtain a FlexiChain{VarName}. This must return a tuple of a DynamicPPL.VarNamedTuple and a NamedTuple:

using DynamicPPL: VarNamedTuple, VarName
FlexiChains.to_vnt_and_stats(::T) = (VarNamedTuple(hello=1.0), (; world=2.0))
sample(M(), S(), 10; chain_type=FlexiChain{VarName})

╭─FlexiChain (10 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter  = 1:10                                                               │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── VarName                                                    │
│  Float64  hello                                                              │
│                                                                              │
│ Extras (1)                                                                   │
│  Float64  world                                                              │
╰──────────────────────────────────────────────────────────────────────────────╯

Docstrings

FlexiChains.to_vnt_and_stats Function

to_vnt_and_stats(transition)::Tuple{VarNamedTuple,NamedTuple}

Convert the first output (i.e. the 'transition') of an AbstractMCMC sampler into a VarNamedTuple mapping parameter names to their values, plus a NamedTuple of any additional statistics.

The VarNamedTuple will be converted into Parameter keys, and the NamedTuple into Extra keys.

If you are writing a custom AbstractMCMC sampler and want to allow users to collect the samples as a FlexiChain{VarName}, i.e.,

sample(...; chain_type=FlexiChain{VarName})

then you should ensure that your method of AbstractMCMC.step returns a transition that can be passed to this method. (The other return value, the state, is not relevant.)

Note that this method is already implemented for DynamicPPL.VarNamedTuple (in which case the stats are empty) as well as DynamicPPL.ParamsWithStats. Thus the easiest solution is to just return one of those types.

source

FlexiChains.to_nt_and_stats Function

to_nt_and_stats(transition)::Tuple{NamedTuple,NamedTuple}

Convert the first output (i.e. the 'transition') of an AbstractMCMC sampler into a NamedTuple mapping parameter names to their values, plus a NamedTuple of any additional statistics.

The first NamedTuple will be converted into Parameter keys, and the second NamedTuple into Extra keys.

If you are writing a custom AbstractMCMC sampler and want to allow users to collect the samples as a FlexiChain{Symbol}, i.e.,

sample(...; chain_type=FlexiChain{Symbol})

then you should ensure that your method of AbstractMCMC.step returns a transition that can be passed to this method. (The other return value, the state, is not relevant.)

source