API

AbstractMCMC defines an interface for sampling Markov chains.

Model

AbstractModel

LogDensityModel <: AbstractMCMC.AbstractModel

Sampler

AbstractSampler

Sampling a single chain

sample(
    rng::Random.AbatractRNG=Random.default_rng(),
    model::AbstractModel,
    sampler::AbstractSampler,
    N_or_isdone;
    kwargs...,
)

isdone(rng, model, sampler, samples, state, iteration; kwargs...)

sample(
    rng::Random.AbstractRNG=Random.default_rng(),
    logdensity,
    sampler::AbstractSampler,
    N_or_isdone;
    kwargs...,
)

Iterator

steps(
    rng::Random.AbstractRNG=Random.default_rng(),
    model::AbstractModel,
    sampler::AbstractSampler;
    kwargs...,
)

julia> struct MyModel <: AbstractMCMC.AbstractModel end

julia> struct MySampler <: AbstractMCMC.AbstractSampler end

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

julia> iterator = steps(MyModel(), MySampler());

julia> collect(Iterators.take(iterator, 10)) == zeros(10)
true

steps(
    rng::Random.AbstractRNG=Random.default_rng(),
    logdensity,
    sampler::AbstractSampler;
    kwargs...,
)

Transducer

Sample(
    rng::Random.AbstractRNG=Random.default_rng(),
    model::AbstractModel,
    sampler::AbstractSampler;
    kwargs...,
)

julia> struct MyModel <: AbstractMCMC.AbstractModel end

julia> struct MySampler <: AbstractMCMC.AbstractSampler end

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

julia> transducer = Sample(MyModel(), MySampler());

julia> collect(transducer(1:10)) == zeros(10)
true

Sample(
    rng::Random.AbstractRNG=Random.default_rng(),
    logdensity,
    sampler::AbstractSampler;
    kwargs...,
)

Sampling multiple chains in parallel

sample(
    rng::Random.AbstractRNG=Random.default_rng(),
    model::AbstractModel,
    sampler::AbstractSampler,
    parallel::AbstractMCMCEnsemble,
    N::Integer,
    nchains::Integer;
    kwargs...,
)

sample(
    rng::Random.AbstractRNG=Random.default_rng(),
    logdensity,
    sampler::AbstractSampler,
    parallel::AbstractMCMCEnsemble,
    N::Integer,
    nchains::Integer;
    kwargs...,
)

Two algorithms are provided for parallel sampling with multiple threads and multiple processes, and one allows for the user to sample multiple chains in serial (no parallelization):

MCMCThreads

MCMCDistributed

MCMCSerial

Common keyword arguments

Common keyword arguments for regular and parallel sampling are:

• progress (default: AbstractMCMC.PROGRESS[] which is true initially): toggles progress logging
• chain_type (default: Any): determines the type of the returned chain
• callback (default: nothing): if callback !== nothing, then callback(rng, model, sampler, sample, iteration) is called after every sampling step, where sample is the most recent sample of the Markov chain and iteration is the current iteration
• num_warmup (default: 0): number of "warm-up" steps to take before the first "regular" step, i.e. number of times to call AbstractMCMC.step_warmup before the first call to AbstractMCMC.step.
• discard_initial (default: num_warmup): number of initial samples that are discarded. Note that if discard_initial < num_warmup, warm-up samples will also be included in the resulting samples.
• thinning (default: 1): factor by which to thin samples.
• initial_state (default: nothing): if initial_state !== nothing, the first call to AbstractMCMC.step is passed initial_state as the state argument.

There is no "official" way for providing initial parameter values yet. However, multiple packages such as EllipticalSliceSampling.jl and AdvancedMH.jl support an initial_params keyword argument for setting the initial values when sampling a single chain. To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, we decided to support initial_params in the default implementations of the ensemble methods:

• initial_params (default: nothing): if initial_params isa AbstractArray, then the ith element of initial_params is used as initial parameters of the ith chain. If one wants to use the same initial parameters x for every chain, one can specify e.g. initial_params = FillArrays.Fill(x, N).

Progress logging can be enabled and disabled globally with AbstractMCMC.setprogress!(progress).

setprogress!(progress::Bool; silent::Bool=false)

Chains

The chain_type keyword argument allows to set the type of the returned chain. A common choice is to return chains of type Chains from MCMCChains.jl.

AbstractMCMC defines the abstract type AbstractChains for Markov chains.

AbstractChains

For chains of this type, AbstractMCMC defines the following two methods.

chainscat(c::AbstractChains...)

chainsstack(c::AbstractVector)

Interacting with states of samplers

To make it a bit easier to interact with some arbitrary sampler state, we encourage implementations of AbstractSampler to implement the following methods:

getparams(state[; kwargs...])

setparams!!(state, params)

These methods can also be useful for implementing samplers which wraps some inner samplers, e.g. a mixture of samplers.

Example: MixtureSampler

In a MixtureSampler we need two things:

• components: collection of samplers.
• weights: collection of weights representing the probability of choosing the corresponding sampler.

struct MixtureSampler{W,C} <: AbstractMCMC.AbstractSampler
    components::C
    weights::W
end

To implement the state, we need to keep track of a couple of things:

• index: the index of the sampler used in this step.
• states: the current states of all the components.

We need to keep track of the states of all components rather than just the state for the sampler we used previously. The reason is that lots of samplers keep track of more than just the previous realizations of the variables, e.g. in AdvancedHMC.jl we keep track of the momentum used, the metric used, etc.

struct MixtureState{S}
    index::Int
    states::S
end

The step for a MixtureSampler is defined by the following generative process

\[\begin{aligned} i &\sim \mathrm{Categorical}(w_1, \dots, w_k) \\ X_t &\sim \mathcal{K}_i(\cdot \mid X_{t - 1}) \end{aligned}\]

where $\mathcal{K}_i$ denotes the i-th kernel/sampler, and $w_i$ denotes the weight/probability of choosing the i-th sampler. AbstractMCMC.getparams and AbstractMCMC.setparams!! comes into play in defining/computing $\mathcal{K}_i(\cdot \mid X_{t - 1})$ since $X_{t - 1}$ could be coming from a different sampler.

If we let state be the current MixtureState, i the current component, and i_prev is the previous component we sampled from, then this translates into the following piece of code:

# Update the corresponding state, i.e. `state.states[i]`, using
# the state and transition from the previous iteration.
state_current = AbstractMCMC.setparams!!(
    state.states[i], 
    AbstractMCMC.getparams(state.states[i_prev]),
)

# Take a `step` for this sampler using the updated state.
transition, state_current = AbstractMCMC.step(
    rng, model, sampler_current, sampler_state;
    kwargs...
)

The full AbstractMCMC.step implementation would then be something like:

function AbstractMCMC.step(rng, model::AbstractMCMC.AbstractModel, sampler::MixtureSampler, state; kwargs...)
    # Sample the component to use in this `step`.
    i = rand(Categorical(sampler.weights))
    sampler_current = sampler.components[i]

    # Update the corresponding state, i.e. `state.states[i]`, using
    # the state and transition from the previous iteration.
    i_prev = state.index
    state_current = AbstractMCMC.setparams!!(  
        state.states[i], 
        AbstractMCMC.getparams(state.states[i_prev]),  
    )

    # Take a `step` for this sampler using the updated state.
    transition, state_current = AbstractMCMC.step(
        rng, model, sampler_current, state_current;
        kwargs...
    )

    # Create the new states.
    # NOTE: Code below will result in `states_new` being a `Vector`.
    # If we wanted to allow usage of alternative containers, e.g. `Tuple`,
    # it would be better to use something like `@set states[i] = state_current`
    # where `@set` is from Setfield.jl.
    states_new = map(1:length(state.states)) do j
        if j == i
            # Replace the i-th state with the new one.
            state_current
        else
            # Otherwise we just carry over the previous ones.
            state.states[j]
        end
    end

    # Create the new `MixtureState`.
    state_new = MixtureState(i, states_new)

    return transition, state_new
end

And for the initial AbstractMCMC.step we have:

function AbstractMCMC.step(rng, model::AbstractMCMC.AbstractModel, sampler::MixtureSampler; kwargs...)
    # Initialize every state.
    transitions_and_states = map(sampler.components) do spl
        AbstractMCMC.step(rng, model, spl; kwargs...)
    end

    # Sample the component to use this `step`.
    i = rand(Categorical(sampler.weights))
    # Extract the corresponding transition.
    transition = first(transitions_and_states[i])
    # Extract states.
    states = map(last, transitions_and_states)
    # Create new `MixtureState`.
    state = MixtureState(i, states)

    return transition, state
end

Suppose we then wanted to use this with some of the packages which implements AbstractMCMC.jl's interface, e.g. AdvancedMH.jl, then we'd simply have to implement getparams and setparams!!.

To use MixtureSampler with two samplers sampler1 and sampler2 from AdvancedMH.jl as components, we'd simply do

sampler = MixtureSampler([sampler1, sampler2], [0.1, 0.9])
transition, state = AbstractMCMC.step(rng, model, sampler)
while ...
    transition, state = AbstractMCMC.step(rng, model, sampler, state)
end