Callbacks

AbstractMCMC provides a unified callback API for monitoring and logging MCMC sampling.

Basic Usage

The mcmc_callback function is the main entry point for creating callbacks:

using AbstractMCMC

struct MyModel <: AbstractMCMC.AbstractModel end

struct MySampler <: AbstractMCMC.AbstractSampler end

function AbstractMCMC.step(rng, ::MyModel, ::MySampler, state=nothing; kwargs...)
   # all samples are zero
   return 0.0, state
end

model, sampler = MyModel(), MySampler()

# Simple callback with a function
cb = mcmc_callback() do rng, model, sampler, transition, state, iteration
    println("Iteration: $iteration")
end

chain = sample(model, sampler, 10; callback=cb)

Combining Multiple Callbacks

Pass multiple callbacks to mcmc_callback to combine them:

cb1 = (args...; kwargs...) -> println("Callback 1")
cb2 = (args...; kwargs...) -> println("Callback 2")

cb = mcmc_callback(cb1, cb2)

You can also add callbacks dynamically using BangBang.push!!:

using BangBang

cb = mcmc_callback(cb1)
cb = push!!(cb, cb2)

TensorBoard Logging

TensorBoard logging requires TensorBoardLogger. Statistics collection also requires OnlineStats.

Basic Logging (No Statistics)

using AbstractMCMC
using TensorBoardLogger

logger = TBLogger("runs/experiment1")
cb = mcmc_callback(logger=logger)

chain = sample(model, sampler, 1000; callback=cb)

Logging with Statistics

To collect running statistics (mean, variance, histograms), load OnlineStats and use the stats argument:

using AbstractMCMC
using TensorBoardLogger
using OnlineStats

logger = TBLogger("runs/experiment1")

# Use default statistics (Mean, Variance, KHist)
cb = mcmc_callback(logger=logger, stats=true)

# Or specify custom statistics
cb = mcmc_callback(
    logger=logger,
    stats=(Mean(), Variance(), KHist(50)),
)

Stats Processing Options

Control how samples are processed before computing statistics with stats_options:

cb = mcmc_callback(
    logger=logger,
    stats=true,
    stats_options=(
        skip=100,    # Skip first 100 samples (burn-in)
        thin=5,      # Use every 5th sample
        window=1000, # Rolling window of 1000 samples
    ),
)

Options merge with defaults, so you only need to specify what you want to change:

# Only change thin, skip and window use defaults (0 and typemax(Int))
cb = mcmc_callback(logger=logger, stats=true, stats_options=(thin=10,))

Name Filtering

Use name_filter to control which parameters and statistics are logged:

cb = mcmc_callback(
    logger=logger,
    name_filter=(
        include=["mu", "sigma"],  # Only log these parameters
        exclude=["_internal"],     # Exclude matching names
        stats=true,                # Include step-level statistics
        extras=true,               # Include extra diagnostics
    ),
)

Complete Example

using AbstractMCMC
using TensorBoardLogger
using OnlineStats

logger = TBLogger("runs/full_example")

cb = mcmc_callback(
    logger=logger,
    stats=true,
    stats_options=(skip=50, thin=2),
    name_filter=(
        exclude=["_internal"],
        extras=true,
        hyperparams=true,
    ),
)

chain = sample(model, sampler, 10000; callback=cb)

Then view in TensorBoard:

tensorboard --logdir=runs/full_example

Navigate to localhost:6006 in your browser to see the dashboard. You'll see real-time plots of your parameter distributions, histograms, and other statistics as sampling progresses.

The Time Series tab provides detailed traces of parameter values throughout the sampling process.

The Scalars tab shows time series of parameter values and statistics over the sampling iterations.

The Distributions tab displays the marginal distributions of each parameter.

The Histograms tab shows the evolution of parameter distributions over time.

API Reference

mcmc_callback(;
    logger,
    stats = nothing,
    stats_options = nothing,
    name_filter = nothing,
)

using TensorBoardLogger
lg = TBLogger("runs/exp")
cb = mcmc_callback(logger=lg)

# With default stats (requires OnlineStats)
using TensorBoardLogger, OnlineStats
lg = TBLogger("runs/exp")
cb = mcmc_callback(logger=lg, stats=true)

ParamsWithStats{P,S,E}

pws = ParamsWithStats(model, sampler, transition, state; params=true, stats=true)
for (name, value) in Base.pairs(pws)
    println("$name: $value")
end

# Re-select to exclude stats:
pws2 = ParamsWithStats(pws; params=true, stats=false)

Default Values

stats_options defaults

name_filter defaults

Implementing Custom Callbacks

Any callable with the following signature can be used as a callback:

function my_callback(rng, model, sampler, transition, state, iteration; kwargs...)
    # Your callback logic here
end

ParamsWithStats

ParamsWithStats is a container for extracting and iterating over MCMC parameters, statistics, and extras.

Basic Usage

# Extract params and stats from state
pws = ParamsWithStats(model, sampler, transition, state; params=true, stats=true)

# Iterate using Base.pairs
for (name, value) in Base.pairs(pws)
    @info name value
end

# Re-select to get only params
pws_params = ParamsWithStats(pws; params=true, stats=false, extras=false)

Overriding for Your Package

To provide meaningful variable names, override the extraction hooks:

# Option 1: Return Vector{<:Real} - default names (θ[1], θ[2], ...) will be used
function AbstractMCMC.getparams(state::MyState)
    return [state.mu, state.sigma]
end

# Option 2: Return named pairs - will be converted to NamedTuple
function AbstractMCMC.getparams(state::MyState)
    return ["μ" => state.mu, "σ" => state.sigma]
end

# Override getstats to return step-level statistics as NamedTuple
function AbstractMCMC.getstats(state::MyState)
    return (lp=state.logp, acceptance_rate=state.accept_rate)
end

The ParamsWithStats constructors normalize all inputs to NamedTuple:

• Vector{<:Real} gets default θ[i] names
• Vector{Pair} is converted to NamedTuple with the provided names
• NamedTuple is used directly

Usage in TensorBoard Callback

The TensorBoard callback uses ParamsWithStats with Base.pairs:

pws = ParamsWithStats(model, sampler, t, state; params=true, stats=true)
for (k, val) in Base.pairs(pws)
    @info "$k" val
end

Internals

Types

MultiCallback

NameFilter(; include=Set{String}(), exclude=Set{String}())

OnlineStats Wrappers

When using statistics, AbstractMCMC provides wrappers that modify how samples are processed:

These are applied automatically via stats_options, but can also be used directly if needed.