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

sample(model, sampler, MCMCThreads(), N, nchains)

MCMCDistributed

sample(model, sampler, MCMCDistributed(), N, nchains)

MCMCSerial

sample(model, sampler, MCMCSerial(), N, nchains)

Common keyword arguments

Common keyword arguments for regular and parallel sampling are:

• progress (default: AbstractMCMC.PROGRESS[] which is true initially): toggles progress logging. See the section on Progress logging below for more details.
• chain_type (default: Any): determines the type of the returned chain
• callback (default: nothing): if callback !== nothing, then callback(rng, model, sampler, sample, iteration; kwargs...) is called after every sampling step, where sample is the most recent sample of the Markov chain and iteration is the current iteration
  • Keyword arguments kwargs... are passed down from the call to sample(...). If you are performing multiple-chain sampling, then kwargs additionally contains a chain_number keyword argument, which runs from 1 to the number of chains. This is not present when performing single-chain sampling.
• 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

Progress logging is controlled in one of two ways:

• by passing the progress keyword argument to the sample(...) function, or
• by globally changing the defaults with AbstractMCMC.setprogress! and AbstractMCMC.setmaxchainsprogress!.

progress keyword argument

For single-chain sampling (i.e., sample([rng,] model, sampler, N)), as well as multiple-chain sampling with MCMCSerial, the progress keyword argument should be a Bool.

For multiple-chain sampling using MCMCThreads, there are several, more detailed, options:

• :overall: create one progress bar for the overall sampling process, which tracks the percentage of samples that have been sampled across all chains
• :perchain: in addition to :overall, also create one progress bar for each individual chain
• :none: do not create any progress bar
• true (the default): same as :overall, i.e. one progress bar for the overall sampling process
• false: same as :none, i.e. no progress bar

Multiple-chain sampling using MCMCDistributed behaves the same as MCMCThreads, except that :perchain is not (yet?) implemented.

Global settings

If you are sampling multiple times and would like to change the default behaviour, you can use this function to control progress logging globally:

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

setprogress! is more general, and applies to all types of sampling (both single- and multiple-chain). It only takes a boolean argument, which switches progress logging on or off. For example, setprogress!(false) will disable all progress logging.

Note that setprogress! cannot be used to set the type of progress bar for multiple-chain sampling. If you want to use :perchain, it has to be set on each individual call to sample.

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

The following two functions exist for converting AbstractChains from and to matrices of samples.

AbstractMCMC.to_samples(::Type{T}, chains::AbstractChains) where {T}

AbstractMCMC.from_samples(::Type{T}, samples::Matrix) where {T<: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([model::AbstractModel, ]state)::Vector{<:Real}

setparams!!([model::AbstractModel, ]state, params)

getparams and setparams!! provide a generic interface for interacting with the parameters of a sampler's state, regardless of how that state is represented internally.

This allows generic code to be written that works with any sampler implementing this interface. For example, a generic ensemble sampler could use getparams to extract the parameters from each of its component samplers' states, and setparams!! to initialize each component sampler with a different set of parameters.

The optional model argument to these functions allows sampler implementations to customize their behavior based on the model being used. For example, some samplers may need to evaluate the log density at new parameter values when setting parameters, which requires access to the model. If access to model is not needed, the sampler only needs to implement the version without the model argument - the default implementations will then call those methods directly.

These methods are particularly useful for implementing samplers which wrap some inner samplers, such as a mixture of samplers. In the next section, we will see how getparams and setparams!! can be used to implement a MixtureSampler.

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