Design/improve/unifty `sample()` interface

[xukai92]

Questions and my current thoughts

• How to pass initial state
  • Support sample(..., init::Array{Union{Missing, Array{Float64,1}},1}=[missing])
  • We flat init, and assign those variables which are not missed to vi[spl]
  • Constructing a correct init is the user's responsibility
• How to pass adaptation configuration
  • Take use of AdvancedHMC.Adaptation
  • sample(..., adapt::Union{Nothing,AbstractAdaptor}=nothing)
• How to pass continuing chain
  • sample(..., resume::Union{Nothing,AbstractChain}=nothing)
  • How do we resume an adapted sampler?
    • With the new interface of AdvancedHMC.Adaptation, the adapted sampler is just the types but not the adapter or any other thing. So if we can make the types in Turing to be 1:1 to those types, we can resume the adapted sampler.

Related issue:

• Provide a way to set initial state of HMC Provide a way to set initial state of HMC #571
• Better adaptation configuration interface Better adaptation configuration interface #491
• Resume interface Resume interface #566

[trappmartin]

Does it make sense that we differentiate between initial value and resume sampling? Why not get rid of the resume interface and only provide a clean way to set an initial value?

[xukai92]

Good idea! Then the only question leaves here is how to resume adapted sampler.

[yebai]

I think we can use keyword arguments here to allow maximum flexibility and minimal sample design. For example:

function sample(rng::AbstractRNG, 
                       ℓ::ModelType, s::SamplerType, 
                       N::Integer; args...
        ) where {ModelType, SamplerType, TransitionType}
    t = Array{TransitionType, 1}(undef, N)
     for i=1:N
        t[i] = step(rng::AbstractRNG, ℓ:: ModelType, s::Sampler, N::Integer; args...)
    end
   # convert `t` into `MCMCChains.Chain`.
end

where args is a keyword argument, some example uses are

• θ₀::AbstractVector = zeros(D)
• adaptor::Adaptor = StanNUTSAdaptor()
• n_chains::Int = 1
• resume::AbstractChain = chain

[yebai]

Related: #723

[yebai]

@cpfiffer Does it make sense moving this sample function to MCMCChains. This design would allow packages such as AdvancedHMC fully functional by depending only on MCMCChains. That is, making both Turing and AdvancedHMC depending on MCMCChains, and removing inter-dependence between them.

[cpfiffer]

I don't see why not. The function looks pretty small, and everything's parametric so there's not really any dependencies to speak of. The JuliaStan people may have some opinions though, so allow me to summon the JuliaStan kingpin.

@goedman, we're working on something of an overhaul of the Turing internals for our 0.7 release and wanted to know what you think about including a function in MCMCChains like the following:

function sample(rng::AbstractRNG, 
    ℓ:: ModelType, s::SamplerType, 
    N::Integer; args...
) where {ModelType, SamplerType, TransitionType}
    t = Array{TransitionType, 1}(undef, N)

    for i=1:N
        t[i] = step(rng::AbstractRNG, ℓ::ModelType, s::SamplerType, N::Integer; args...)
    end
    
    return Chains(t ...)
end

It's designed such that AdvancedHMC and Turing can work independently of one another. I'd view it as basically another MCMCChains constructor. Do you have any issues with including this in MCMCChains?

[goedman]

That's perfectly ok.

[goedman]

I've actually started to look at the AHMC sampler would love to spend some time with it next week and try it out! If the performance is reasonably close to DynamicHMC that would be a huge leap forward!

[cpfiffer]

@yebai Does this interface need some additional components? I've added all the bits to a draft of MCMCChains that includes the addition of AbstractSampler, AbstractModel, and AbstractTransition types which I could see Turing inheriting from. I can take this out, I'm not very convinced as to their necessity, but it does unify the "interface" component of all this.

I'm not sure on what AbstractTransition should look like, though, and I couldn't find an example of what this might be. The function above also doesn't have a mechanism for the type to actually make it into the function. Should AbstractModel have some parametric type AbstractModel{TransitionType}? This comment seems to indicate that the TransitionType is at the sampler level.

Additionally, many of the samplers have some initial setup of some kind that isn't caught well by the "iterate through each step" methodology above. I added some "empty functions" that specific samplers can overload if they need to. It would be really cool if this could catch the general cases for all the samplers.

For example:

function sample(
    rng::AbstractRNG,
    ℓ::ModelType, 
    s::SamplerType,
    N::Integer; 
    kwargs...
) where {ModelType<:AbstractModel,
    SamplerType<:AbstractSampler,
    TransitionType<:AbstractTransition} #Not sure how to get TransitionType into the function
    t = Array{TransitionType, 1}(undef, N)

    # Perform any necessary setup.
    sample_init!(rng, ℓ, s, N; kwargs...)

    # Step through the sampler.
    for i=1:N
        t[i] = step(rng::AbstractRNG, ℓ::ModelType, s::SamplerType, N::Integer; kwargs...)
    end

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

   # Placeholder while I figure out what t looks like.
    return Chains(t)
end

In summary, this is what I see the sampling interface looking like:

• Samplers inherit from AbstractSampler or some subtype of AbstractSampler, like Hamiltonian<:AbstractSampler
• Models inherit from the AbstractModel type. This won't be too hard on the Turing side, as

  struct Model{pvars, dvars, F, TData, TDefaults}

  would just become

  struct Model{pvars, dvars, F, TData, TDefaults} <: AbstractModel

• Samplers that need setup or wrap-up code overload their own versions of sample_init!(rng, ℓ, s, N; kwargs...) and sample_end!(rng, ℓ, s, N; kwargs...).
• I have no idea what to do with TransitionType or AbstractTransition.

I also have wrapper functions for the above that allow for the sampling of multiple chains, and psample which will run multiple chains on parallel processes if any workers are available, but I can't really test them well without the whole interface working.

Comments, pointers, and corrections welcomed.

[yebai]

Some desired refactoring goals for Sampler APIs

1. Implement a shared sample(model, alg, ...) for all samplers (related to https://github.com/TuringLang/Turing.jl/issues/634#issuecomment-471339521, Pass number of parallel chains as an argument #582, Add example illustrating how to run multiple chains. #698, Better adaptation configuration interface #491, Add additional keyword arguments to the sample function #482). This would substantially simply sampling algorithm code since their current implementations contain a customised sample with loads of code copied and pasted from the HMC sampler.

2. Remove Sampler.info and introduce a Transitiontype (generalising current Turing.Sample type) to store all MCMC running time information (related Logging in samplers / diagnosis statistics #599). This would remove the need to create a Sampler wrapper for each Gibbs component algorithms.
   

   Turing.jl/src/utilities/io.jl

   Lines 9 to 12 in 6c0f4b8

   	mutable struct Sample
   	weight :: Float64 # particle weight
   	value :: Dict{Symbol,Any}
   	end

3. Remove Sampler wrapper for Gibbs component algorithms: Gibbs.algs::Sampler ==> Gibbs.algs:InferenceAlgorithm. (related: Introduce MCMCSampler type #602, [Gibbs] Increase performance when subsamplers have several iterations #378)
   

   Turing.jl/src/inference/gibbs.jl

   Lines 44 to 52 in 6c0f4b8

   	for i in 1:n_samplers
   	sub_alg = alg.algs[i]
   	if isa(sub_alg, GibbsComponent)
   	samplers[i] = Sampler(typeof(sub_alg)(sub_alg, i), model)
   	else
   	@error("[Gibbs] unsupport base sampling algorithm $alg")
   	end
   	space = union(space, sub_alg.space)
   	end

Related https://github.com/TuringLang/Turing.jl/issues/634 #602 #582 #599 #689

Here are some relevant sample implementations in the current code base:

HMC, SGLD and SGHMC:

• Turing.jl/src/inference/hmc.jl

  Lines 214 to 224 in d2a987d

  	function sample(
  	model::Model,
  	alg::Hamiltonian;
  	save_state=false, # flag for state saving
  	resume_from=nothing, # chain to continue
  	reuse_spl_n=0, # flag for spl re-using
  	adaptor = AHMCAdaptor(alg),
  	init_theta::Union{Nothing,Array{<:Any,1}}=nothing,
  	rng::AbstractRNG=GLOBAL_RNG,
  	kwargs...
  	)

Gibbs

• Turing.jl/src/inference/gibbs.jl

  Lines 70 to 76 in d2a987d

  	function sample(
  	model::Model,
  	alg::Gibbs;
  	save_state=false, # flag for state saving
  	resume_from=nothing, # chain to continue
  	reuse_spl_n=0 # flag for spl re-using
  	)

PGibbs

• Turing.jl/src/inference/pgibbs.jl

  Lines 82 to 87 in d2a987d

  	function sample( model::Model,
  	alg::PG;
  	save_state=false, # flag for state saving
  	resume_from=nothing, # chain to continue
  	reuse_spl_n=0 # flag for spl re-using
  	)

SMC

• Turing.jl/src/inference/smc.jl

  Lines 71 to 72 in d2a987d

  	## wrapper for smc: run the sampler, collect results.
  	function sample(model::Model, alg::SMC)

Stan

• Turing.jl/src/utilities/stan-interface.jl

  Lines 34 to 41 in d2a987d

  	function sample(mf::T,
  	num_samples::Int,
  	num_warmup::Int,
  	save_warmup::Bool,
  	thin::Int,
  	adapt::CmdStan.Adapt,
  	alg::CmdStan.Hmc
  	) where T