AbstractChains interface

[penelopeysm]

AbstractMCMC defines the AbstractChains type:

AbstractMCMC.jl/src/AbstractMCMC.jl

Lines 25 to 31 in ae9760d

	"""
	AbstractChains
	`AbstractChains` is an abstract type for an object that stores
	parameter samples generated through a MCMC process.
	"""
	abstract type AbstractChains end

plus two methods on it:

AbstractMCMC.jl/src/interface.jl

Lines 1 to 20 in ae9760d

	"""
	chainscat(c::AbstractChains...)
	Concatenate multiple chains.
	By default, the chains are concatenated along the third dimension by calling
	`cat(c...; dims=3)`.
	"""
	chainscat(c::AbstractChains...) = cat(c...; dims=3)
	"""
	chainsstack(c::AbstractVector)
	Stack chains in `c`.
	By default, the vector of chains is returned unmodified. If `eltype(c) <: AbstractChains`,
	then `reduce(chainscat, c)` is called.
	"""
	chainsstack(c) = c
	chainsstack(c::AbstractVector{<:AbstractChains}) = reduce(chainscat, c)

This skeletal interface was alright when MCMCChains.Chains was the only concrete type that implemented it. However, it seems clear that there is some benefit to be obtained by fleshing out the interface. This would probably enable a fair amount of code reuse.

In particular, I would like to see something like the following (modulo names, and this list is hardly exhaustive!!)

Base.keys(c::AbstractChains)
Base.values(c::AbstractChains)
values_at(c::AbstractChains, iter, chain) # this is needed for things like DynamicPPL returned/predict
Base.pairs(c::AbstractChains)
Base.size(c::AbstractChains)
Base.vcat(c1, c2) # concatenates iters
Base.hcat(c1, c2) # concatenates chains
Base.merge(c1, c2) # merges parameters
get_data(c::AbstractChains, key) # --> matrix of (iters, chains)
subset(c::AbstractChains, keys) # --> returns the same AbstractChains but with only the selected keys
get_iter_indices(c::AbstractChains)
get_chain_indices(c::AbstractChains)
get_sampling_time(c::AbstractChains)
get_saved_sampler_state(c::AbstractChains)
parameters(c::AbstractChains) # Could be StatsBase.params, but I don't like the idea of overloading that
non_parameter_keys(c::AbstractChains)
summarize(c::AbstractChains) # I've been using StatsBase.summarystats for this
Plots.plot(c::AbstractChains[, key_or_keys]; kwargs...) # in an extension
Makie.plot(c::AbstractChains[, key_or_keys]; kwargs...) # in an extension

There can be a default implementation of Base.getindex but I think for the most part the implementation of Base.getindex should be left to the individual types to implement, since different types may want to present different indexing interfaces.

Finally, I would say it's quite likely that this interface should be split into a different package altogether since consumers of AbstractChains may not want to take on AbstractMCMC.jl as a dependency.

Note that this is quite a nontrivial task because not only does the interface need to be figured out, one would also have to work on MCMCChains to implement the interface for it, before any real benefits are obtained. Just declaring an interface alone but not making use of it in MCMCChains won't help with code duplication.

[penelopeysm]

https://github.com/TuringLang/AbstractChains.jl