Implicit dependence on MCMCChains

[torfjelde]

Currently there are certain methods which implicitly make assumptions about the functionality of AbstractMCMC.AbstractChains, e.g. indexing, etc. which are really only present in the particular implementation (though I believe it's also the only implementation) of AbstractChains provided by MCMCChains.Chains.
Thus, we're implicitly depending on MCMCChains.Chains without making it explicit.

In addition, it ends up tying the testing quite heavily to Turing.jl, which is non-ideal.

Possible solutions

1. Converting chains to "dummy" Vector{<:VarInfo}

As breifly discussed in #166 (comment), it seems like one feature that would be very helpful in moving away from this is to provide a method which maps (::Model, ::AbstractChains) -> ::Vector{<:VarInfo}, and vice-versa. Then we could stick to working with VarInfo in DynamicPPL.jl, which would help alleviate the above mentioned issues.

It could for example look something like:

import DynamicPPL: Model, VarInfo, getdist
# Adapted from `Turing.Inference.transitions_from_chain`
function varinfo_from_chain(
    model::Turing.Model,
    chain::MCMCChains.Chains;
    sampler = DynamicPPL.SampleFromPrior()
)
    vi = Turing.VarInfo(model)

    vis = map(1:length(chain)) do i
        c = chain[i]
        md = vi.metadata
        for v in keys(md)
            for vn in md[v].vns
                vn_sym = Symbol(vn)

                # Cannot use `vn_sym` to index in the chain
                # so we have to extract the corresponding "linear"
                # indices and use those.
                # `ks` is empty if `vn_sym` not in `c`.
                ks = MCMCChains.namesingroup(c, vn_sym)

                if !isempty(ks)
                    # 1st dimension is of size 1 since `c`
                    # only contains a single sample, and the
                    # last dimension is of size 1 since
                    # we're assuming we're working with a single chain.
                    val = copy(vec(c[ks].value))
                    DynamicPPL.setval!(vi, val, vn)
                    DynamicPPL.settrans!(vi, false, vn)
                else
                    error("$vn not present in chain but is required by the model")
                end
            end
        end
        new_vi = VarInfo(vi, sampler, vi[sampler])
        setlogp!(new_vi, first(chain[i][:lp])) # Is there a better way?

        return new_vi
    end

    return vis
end

which works fairly well (it does what it's supposed to + haven't found any issues yet).

Only downside with the above is that it requires /one/ additional execution of model everytime we convert.

2. ???

[cpfiffer]

We could actually probably just write this an an extension of bundle_samples inside Turing.jl and treat it as a chain type, so you could just call

chain = sample(model, spl, 1000, chain_type=DynamicPPL.VarInfo)

Of course, this would require us to use VarInfo internally instead of Transition, which might not actually be that bad an idea.

[devmotion]

I think there are (at least) two points here:

• Implementing DynamicPPL-specific functionality with VarInfo and in particular VarName seems useful since it is the most direct way for getting information about variables and samples (that does not exclude the possibility to support other inputs as well, by either adding specific implementations for them or providing appropriate conversion utilities)
• If AbstractMCMC would contain a simple interface for chains (and not only the type AbstractChains), most of the functionality could be implemented in a very generic way. E.g., generated_quantities could be generalized if there is a generic way to iterate over samples (some AbstractMCMC.eachsample similar to eachcol/eachrow) and over each variable in each sample (e.g., some AbstractMCMC.eachparam) that yields name and value. AbstractMCMC could define default implementations for AbstractVector chains and NamedTuple samples, similar to the defaults in Tables such that chains of type Vector{<:NamedTuple} would be supported automatically. MCMCChains would implement AbstractMCMC.eachsample, and if it returns a NamedTuple (should we do that?), everything would work automatically, without explicitly converting between Chains and VarInfo objects. DynamicPPL could e.g. define AbstractMCMC.eachparam(::VarInfo) that iterates over all variables in a VarInfo etc.

[cpfiffer]

I actually kind of like the Tables-style interface you proposed. As to returning a NamedTuple from eachsample, I'm not sure -- my preference would be to return the Transition given by the sampler, but Chains don't ever have access to that information since MCMCChains and others implicitly assume linearized inputs.

[torfjelde]

Implementing DynamicPPL-specific functionality with VarInfo and in particular VarName seems useful since it is the most direct way for getting information about variables and samples (that does not exclude the possibility to support other inputs as well, by either adding specific implementations for them or providing appropriate conversion utilities)

100% agree. But I think at the moment it's a bit unclear where these conversion utilities ought to go. We can't put them in DynamicPPL (well, some, but not direct ones, e.g. from Vector{<:VarInfo} to MCMCChains.Chains, without Requires.jl which I don't think is a path we want to go down), nor in MCMCChains. So just shove it all in Turing? Then that means that in practice, if you wanted to for example use DynamicPPL + MCMCChains in something not-Turing-related, you'd still have to depend on Turing.

If AbstractMCMC would contain a simple interface for chains (and not only the type AbstractChains), most of the functionality could be implemented in a very generic way.

This is the sort of thing I'd like to see. We need a bit more assumptions based on the chains baked into AbstractMCMC.jl regarding the "chain" structure (broader than just subtype of AbstractChains), and methods like you suggest would go far in doing that without being unreasonable.

As to returning a NamedTuple from eachsample, I'm not sure -- my preference would be to return the Transition given by the sampler

Is there anything stopping us from doing both? Also, the annoying thing about the Transition stuff is that there doesn't seem to be any common interface? I understand why it's difficult to do, since samplers have different properties you'd like to return.

At the same time I agree that maybe NamedTuple is too simple.

Does it not seem like it would be a good idea to have something "similar" to VarInfo? That is, the parameters and logp are always present, but then there's a metadata side which can be typed differently depending on the sampler. And then there's a set of methods that we expect to be implemented, e.g. params for returning the parameters in some particular form (could also be specified on a case-by-case basis, assuming certain functionality is available). Though I'd argue something like this is further down the priority queue than the above suggested changes.

[cpfiffer]

I mean, why can't we just use VarInfo? From what I understand it's not that much bigger than a NamedTuple, and it has all the metadata you'd need.

[cpfiffer]

And as to the Vector{<:VarInfo}, we can put that in MCMCChains.jl, since we already have some Requires stuff there to add bundle_samples extensions.

[ElOceanografo]

Some kind of NamedTuple-ish structure or Tables-like interface for the chains would be great IMO. One of my common frustrations is fitting a model with array-valued parameters, and then having to excavate and reconstruct them out of an MCMCChain object. I'd love to be able to write e.g.

chn = sample(my_model, NUTS(), 100)
for samp in chn
    my_function(samp.my_array_parameter)
end

[yebai]

MCMCChains is now a package extension.