Sampling Output Formats

AbstractMCMC.sample builds its result from the chain_type keyword. JuliaBUGS supports three formats, all carrying the same draws: the model parameters, the forward-sampled generated quantities, and the sampler's own statistics.

using JuliaBUGS
using AbstractMCMC
using MCMCChains

model_def = @bugs begin
    mu ~ dnorm(0, 1)
    for i in 1:N
        y[i] ~ dnorm(mu, 1)
    end
    doubled = 2 * mu
end
model = model_def((; N = 3, y = [1.2, 0.8, 1.5]))

BUGSModel (parameters are in transformed (unconstrained) space, with dimension 1):

  Model parameters:
    mu

  Variable sizes and types:
    N: type = Int64
    y: size = (3,), type = Vector{Float64}
    mu: type = Float64
    doubled: type = Float64

Vector{ParamsWithStats}

This is the default: sample without a chain_type returns one AbstractMCMC.ParamsWithStats per draw, the same as chain_type = Vector{ParamsWithStats}. It is the most general of the three and the only one that needs no extra package, since AbstractMCMC is already a dependency. It is not the most compact: a chain stores one array, whereas this repeats the VarName keys for every draw.

draws = AbstractMCMC.sample(model, JuliaBUGS.IndependentMH(), 500; progress = false)
draws[1]

AbstractMCMC.ParamsWithStats{OrderedCollections.OrderedDict{VarName, Any}, @NamedTuple{lp::Float64}, @NamedTuple{}}(OrderedCollections.OrderedDict{VarName, Any}(mu => 0.6812737324835343, doubled => 1.3625474649670686), (lp = -4.384563866270414,), NamedTuple())

Each entry has two fields. params is an OrderedDict keyed by VarName, so array-valued variables are stored whole rather than split into scalar columns:

draws[1].params[@varname(mu)]

0.6812737324835343

stats is a NamedTuple of the statistics the sampler reported for that draw. Only what the sampler actually produced is recorded, so a statistic missing from one draw is absent from its NamedTuple rather than padded. Base.pairs walks parameters and statistics at once:

using AdvancedHMC, ADTypes, ReverseDiff, DifferentiationInterface

ad_model = compile(model_def, (; N = 3, y = [1.2, 0.8, 1.5]); adtype = AutoReverseDiff())
nuts_draws = AbstractMCMC.sample(
    ad_model,
    NUTS(0.8),
    200;
    chain_type = Vector{ParamsWithStats},
    n_adapts = 100,
    discard_initial = 100,
    progress = false,
)
collect(Base.pairs(nuts_draws[1]))

15-element Vector{Pair}:
 Pair{VarName, Any}(mu, 1.1212054893215877)
 Pair{VarName, Any}(doubled, 2.2424109786431754)
 Pair{Symbol, Real}(:lp, -4.430738418762856)
 Pair{Symbol, Real}(:n_steps, 1)
 Pair{Symbol, Real}(:is_accept, true)
 Pair{Symbol, Real}(:acceptance_rate, 1.0)
 Pair{Symbol, Real}(:log_density, -4.430738418762856)
 Pair{Symbol, Real}(:hamiltonian_energy, 4.6468210476116285)
 Pair{Symbol, Real}(:hamiltonian_energy_error, -0.06578091218292581)
 Pair{Symbol, Real}(:max_hamiltonian_energy_error, -0.06578091218292581)
 Pair{Symbol, Real}(:tree_depth, 1)
 Pair{Symbol, Real}(:numerical_error, false)
 Pair{Symbol, Real}(:step_size, 0.48318055619654676)
 Pair{Symbol, Real}(:nom_step_size, 0.48318055619654676)
 Pair{Symbol, Real}(:is_adapt, false)

A plain vector carries no iteration indices, so discard_initial and thinning are not recorded in it. Pass them back when converting (see below) if you need the chain to report the iteration numbers the run actually used.

Multiple chains come back as one vector per chain:

chains = AbstractMCMC.sample(model, JuliaBUGS.IndependentMH(), MCMCThreads(), 500, 4)
length(chains)     # 4
length(chains[1])  # 500

Samplers JuliaBUGS has no transition_params_and_stats method for keep AbstractMCMC's behaviour and hand back their raw transitions.

MCMCChains.Chains

chain_type = Chains returns an MCMCChains.Chains (requires using MCMCChains). Array-valued variables and statistics are flattened into one scalar column per element, and the statistics go into the internals section.

FlexiChains.VNChain

chain_type = VNChain returns a FlexiChains.FlexiChain{VarName} (requires using FlexiChains). Draws are keyed by VarName with array-valued variables kept whole, and statistics are stored as FlexiChains.Extra entries.

Converting between formats

A vector of ParamsWithStats converts into either chain type with AbstractMCMC.from_samples. It takes a matrix of draws, iterations down the rows and chains across the columns, so a single run needs a reshape:

chain = AbstractMCMC.from_samples(Chains, reshape(draws, :, 1))
summarystats(chain)

Summary Statistics

  parameters      mean       std      mcse   ess_bulk   ess_tail      rhat   e ⋯
      Symbol   Float64   Float64   Float64    Float64    Float64   Float64     ⋯

          mu    0.6840    0.4605    0.0352   169.5701   164.8054    1.0043     ⋯
     doubled    1.3679    0.9210    0.0703   169.5701   164.8054    1.0043     ⋯

                                                                1 column omitted

start and thin restore the iteration numbering a run used, since the vector does not carry it: for discard_initial = n pass start = n + 1.

chain = AbstractMCMC.from_samples(Chains, reshape(draws, :, 1); start = 101, thin = 2)

Several chains go in as columns:

chain = AbstractMCMC.from_samples(Chains, reduce(hcat, chains))

using FlexiChains
chain = AbstractMCMC.from_samples(VNChain, reshape(draws, :, 1))

For ArviZ.jl, go through VNChain: FlexiChains ships an InferenceObjects extension that keeps array-valued variables whole and maps the HMC statistic names onto ArviZ's conventions (hamiltonian_energy to energy, numerical_error to diverging). Converting via Chains works too but flattens arrays into x[1], x[2] and loses that mapping.

Reproducing generated quantities

Generated quantities are forward-sampled when the draws are laid out, from a stream seeded by the draws themselves and the chain number. A seeded sampling run therefore reconstructs identically, and parallel chains get independent streams, with nothing extra to pass.

What a callback sees

AbstractMCMC.ParamsWithStats(model, sampler, transition, state) inside an mcmc_callback reports the parameters the sampler moves, keyed by VarName, and the same statistics as the sampling output, for every sampler. It holds only what the sampler moves: generated quantities, and under auto-marginalization the marginalized discrete latents, are reconstructed once at the end of a run, so they appear only in the sampling output.

For samplers whose transitions do not already carry the model's variables, which is every sampler except Gibbs and IndependentMH, naming a draw costs one model evaluation per iteration. That is worth knowing when a callback runs on every step of a long chain.

Supporting a new sampler

A sampler works with all three formats once it implements a single method that unpacks its transitions:

function JuliaBUGS.transition_params_and_stats(::BUGSModel, ::MySampler, t::MyTransition)
    return t.params, (; lp = t.lp)
end

params is the flat parameter vector for that draw, ordered as LogDensityProblems.logdensity expects it, and stats is whatever the sampler reported. Report only what it actually produced: leave a statistic out of the NamedTuple rather than padding it, and array-valued statistics are fine. Chains takes the union of the keys across draws and flattens arrays itself.

This one method also drives what a callback sees, so implementing it is enough for every output format.