Wrapping an external sampler: PolyChord

[aplavin]

Recently we had a short discussion on discourse (https://discourse.julialang.org/t/mcmc-landscape/25654/83 and around) about using an external nested sampling implementation (e.g. PolyChord) with Turing. I made a thin wrapper (https://github.com/aplavin/Polychord.jl) around the underlying C/Fortran library half a year ago that works fine, but it would be nice to benefit from Turing model definition language. However, I couldn't find a simple way to add this a a Turing sampler - looking at existing samplers it looks like the control flow has somewhat "inverted" direction in Turing and in PolyChord library, at least regarding priors. The workflow for the latter looks like this:

# `prior` converts from uniform cube to model parameters
# e.g. if we have two parameters, one with Normal(0, 1) prior and the other with Normal(5, 2):
prior(cube::Vector{Float})::Vector{Float} = invcdf.([Normal(0, 1), Normal(5, 2)], cube)

# computes likelihood for given model parameters
loglike(params::Vector{Float})::Float = ...

run_polychord(prior, loglike, options)

Any suggestions on how to wrap it as a Turing sampler?

[trappmartin]

I think this should now be possible to implement using the new interface for inference algorithms done by @cpfiffer.

[williamjameshandley]

Hi all. I met with @yebai and @xichenamoy today in the Astro department (with Mike Hobson and Anthony Lasenby), and after that conversation I would be happy to try to get this working with PolyChord. Has anybody interfaced an external sampler yet under @cpfiffer 's new scheme, and if so, could you point me to the part of the code that does?

[cpfiffer]

No, not yet. You would probably be the first. There is a guide I'm working on that shows how to implement MH from the ground up, but at this point it's a little ways off. For now, I can give you a brief overview of how things work. It's a very brief overview and we're still trying to figure out the best way to make development a pleasurable experience, so please let me know if you're unclear anywhere.

The whole interface is built on you overloading various functions as needed. In the most general non-Turing way, you define some AbstractSampler, a Model type, an AbstractTransition type, and then sample_init!, sample_end!, and step! functions as needed.

For the following example, lets pretend you've made a subtype of Turing's InferenceAlgorithm called PolyChord. When you call sample(model, PolyChord(), 1000), the following function is run:

function sample(
    rng::AbstractRNG,
    ℓ::ModelType,
    s::SamplerType,
    N::Integer;
    progress::Bool=true,
    kwargs...
) where {ModelType<:Sampleable, SamplerType<:AbstractSampler}
    # Perform any necessary setup.
    sample_init!(rng, ℓ, s, N; kwargs...)

    # Preallocate the TransitionType vector.
    ts = transitions_init(rng, ℓ, s, N; kwargs...)

    # Add a progress meter.
    cb = progress ? _generate_callback(rng, ℓ, s, N; kwargs...) : nothing

    # Step through the sampler.
    for i=1:N
        if i == 1
            ts[i] = step!(rng, ℓ, s, N; iteration=i, kwargs...)
        else
            ts[i] = step!(rng, ℓ, s, N, ts[i-1]; iteration=i, kwargs...)
        end

        # Run a callback function.
        progress && callback(rng, ℓ, s, N, i, ts[i], cb; kwargs...)
    end

    # Wrap up the sampler, if necessary.
    sample_end!(rng, ℓ, s, N, ts; kwargs...)

    return Chains(rng, ℓ, s, N, ts; kwargs...)
end

This function is the most general possible. It's intended that you'll overload any function in the body of this function (except for _generate_callback, which is strictly internal and shouldn't be overloaded). Everything else is fair game to overload.

There's a lot here, so let's start with a general description of each element.

Structs and types

AbstractSampler

This is some sampler object that describes how inference is done, i.e. how you draw samples in each iteration. In Turing's case, we provide a big general use object called Sampler:

Turing.jl/src/Turing.jl

Line 114 in 4a10e5b

mutable struct Sampler{T, S<:AbstractSamplerState} <: AbstractSampler

Sampler stores

• an inference algorithm like NUTS or MH. Your adaptation would store Polychord or something, whatever you want to call it.
• info, which is mostly a mechanical thing you don't need to care about for now.
• selector, which you also don't need to care about for now.
• state, which accepts some subtype of AbstractSamplerState. This is a custom built struct that maintains state information between samples. You can see an example of this for HMC here.

Here's an example of the state for IS:

mutable struct ISState{V<:VarInfo, F<:AbstractFloat} <: AbstractSamplerState
    vi                 ::  V
    final_logevidence  ::  F

# VarInfo(model) makes a typed VarInfo -- this is important for performance reasons.
ISState(model::Model) = ISState(VarInfo(model), 0.0)

Whatever your state implementation is, make sure it has a vi field that is a subtype of VarInfo. This is how Turing tracks all the variables in your models, and it's generally the most important thing to include to make everything work right.

Model

In this case you don't really need to do anything with the model -- Turing handles all that stuff.

Transition

In interface-lingo, an AbstractTransition type just stores all the information you want to return to the user at the end of sampling. The obvious use here is just the parameter samples. Turing has a default type called Transition which is good enough for most basic samplers. You can get crazy and add some other features if you need to return them, but we can cross that bridge later.
As long as your state variable has a field called vi which stores a VarInfo, you can just call return transition(spl) at the end of each step! function.

You'll need to implement a function transition_type which returns the type of AbstractTransition returned by a given sampler. Here's an example for a particle sampler. You might be able to use transition_type(spl::Sampler{<:PolyChord}) = Transition if you're using the default transition type.

Functions

Sampler

This is not really a general interface thing, but a specific-to-Turing feature. Turing uses InferenceAlgorithms and then stores them in a Sampler type, so you need to make a function with the signature Sampler(alg::PolyChord, model::Model, s::Selector).

IS is the simplest sampler, so you might use it's sampler implementation as an example:

function Sampler(alg::IS, model::Model, s::Selector)
    info = Dict{Symbol, Any}()
    state = ISState(model)
    return Sampler(alg, info, s, state)
end

sample_init!

This is where you do any setup necessary for your sampler. You might generate a particle field and stick it in your sampler.state field, or you might find a good starting position, etc. Not all samplers need this. If you don't you don't need to overload anything.

By default, all Turing samplers have this function called:

function Interface.sample_init!(
    ::AbstractRNG,
    model::Model,
    spl::Sampler,
    N::Integer;
    kwargs...
)
    # Resume the sampler.
    set_resume!(spl; kwargs...)

    # Set the parameters to a starting value.
    initialize_parameters!(spl; kwargs...)
end

All this does is resume a chain if an existing chain was given. The second line checks to see if the user has passed in a starting vector and initializes the sampler to those values. You may or may not need to overload this for PolyChord, I'm not sure. IS doesn't use any special overloading of this function.

step!

The step! function is where all the meat of your sampling is. It is repeated N times, so if you called sample(model, PolyChord(), 1000), step! would run 1000 times and store the result in a vector of whatever your transition type is.

In this case, IS is again the simplest:

function step!(
    ::AbstractRNG,
    model::Model,
    spl::Sampler{<:IS},
    ::Integer;
    kwargs...
)
    empty!(spl.state.vi)
    model(spl.state.vi, spl)

    return Transition(spl)
end

First, this empties out the VarInfo, which is occasionally useful for stateless samplers like IS or SMC. Second, it runs the model you've given it, updates the log density. Finally, it returns the parameter draws by construction a new Transition.

sample_end!

This isn't really used much, but it can be useful if you want to do a finally transformation of your Transitions or save your sampler to disk. I wouldn't worry about this one for now.

Chains

You also shouldn't need to make a new Chains function, since the default built into Turing can handle all the Transition types already built into Turing. If you make a new one that has more specific sampling information, I can help you modify your Transition to get it to work seamlessly with Turing's internals, but that's another bridge we can cross later.

[cpfiffer]

I forgot, you need to also build observe and assume statements. I think probably in general it's a good idea to peruse the IS code to get a really basic introduction to how things flow.

[xukai92]

IS might not be typical because it's not a MCMC sampler. I wonder if we can clean up the MH code a bit and make that as an example.

[yebai]

@cpfiffer Good idea to improve our MH sampler, in terms of modularity, better interface, and advanced features (e.g. adaptive variants).

Some recent related discussions:

• Adaptive Metropolis Algorithm (RAM)
  • https://discourse.julialang.org/t/mcmc-landscape/25654
  • https://github.com/JuliaStats/Klara.jl/blob/master/src/samplers/RAM.jl
• MH-sampler in BAT.jl
  • https://discourse.julialang.org/t/ann-dynamichmc-2-0/28314/22
  • https://github.com/bat/BAT.jl/tree/master/src/samplers/mcmc
• Also, https://github.com/anthofflab/RobustAdaptiveMetropolisSampler.jl

Related issue: https://github.com/TuringLang/Turing.jl/issues/889

I'm happy to help along the process!

[cpfiffer]

It looks like David Anthoff implemented RAM last month. This also looks pretty easy to migrate over to the interface.

[yebai]

Yes; I added it here a few days ago. I feel we can have a really flexible/advanced MH sampler in Turing, by falling back to different MH backends (e.g. Ensemble, RAM, vanilla MH).

[yebai]

It should be significantly easier to get the prior probability of parameters, and likelihood probabilities separately after #989 is merged. This might pave the way for integrating/implementing nested sampling algorithms.

[yebai]

Another implementation of nested sampling: https://github.com/joshspeagle/dynesty

Speagle, J. S. (2019) ‘dynesty: A Dynamic Nested Sampling Package for Estimating Bayesian Posteriors and Evidences’, arXiv [astro-ph.IM]. Available at: http://arxiv.org/abs/1904.02180.

[JohannesBuchner]

Even simpler to implement than MultiNest (dynesty and nestle implement it in Python) and parameter-free is MLFriends. It is illustrated and explained a bit at https://johannesbuchner.github.io/UltraNest/method.html

[yebai]

See https://github.com/TuringLang/NestedSamplers.jl