Integrate Turing.jl probablistic programming

ablaom commented 5 years ago

https://github.com/alan-turing-institute/MLJ.jl/issues/47#issue-402035942

I think Turing.jl developer @trappmartin has indicated an interest in this and this will be explored at upcoming Sprint

@tlienart whose is joining the Turing Institute team soon, may also be interested in helping out.

trappmartin commented 5 years ago

Correct, I’m looking forward to talk about interfacing Turing and MLJ while I’m in London.

fkiraly commented 5 years ago

agreed, it's the right time now. Suggest we start with vanilla Bayesian linear regression, conjugate priors & Gaussian likelihood.

trappmartin commented 5 years ago

Here are a few thoughts on integrating Turing into MLJ for simple models (without local variables). Using more complex models is more involved as we currently do not assume that a Turing program can always be represented by a graphical model to allow arbitrary probabilistic programs. We will likely support the extraction of dependency graphs for programs that allow a graphical model representation in the future.

We would need to:

Implement a mechanism to store to log pdf values for each observation (training or test) to be able to (a) apply model selection and (b) use a Turing model for predictions. See issue TuringLang/Turing.jl#817
Provide an interface to compute the log pdf values, related to PR 750
Implement a wrapper for Turing. Here is an example implementation that could be used once 1 and 2 are resolved:

import StatsFuns: logaddexp

struct TuringModel
   mf::Function
   algo::Turing.InferenceAlgorithm
end

function MLJBase.fit(m::TuringModel, x)
    mf = m.mf
    algo = m.algo

    model = mf(x)
    return Turing.sample(mf(x), algo)
end

function MLJBase.predict(m::TuringModel, fitresult::MCMCChains.Chains, xnew; average = false)
    model = m.mf(xnew)

    vi = Turing.VarInfo()
    model(vi, Turing.SampleFromPrior()) # to obtain VarNames
    vsyms = Turing.RandomVariables.syms(vi)

    lp = ones(size(xnew, 1)) * -Inf

    if average
        vi = Turing.VarInfo()
        y = Vector()
        for s in vsyms
            y = vcat(y, vec( mean(Array(fitresult[:,s,1]), dims=1) ) )
        end
        lp = Turing.logpdf(model, vi, aggregate = false))
    else
        for it = 1:size(fitresult, 1)
            vi = Turing.VarInfo()
            y = Vector()
            for s in vsyms
                y = vcat(y, vec( Array(fitresult[it,s,1]) ) )
            end

            # This is currently not available and the actual syntax might be different.
            lp_ = Turing.logpdf(model, vi, aggregate = false)) - log(size(fitresult,1))
            for i in 1:size(x,1)
                lp[i] = logaddexp(lp[i], lp_[i])
            end
        end
    end
    return lp
end

Additionally, we could think about providing an interface in Turing that allows the user to do:

Turing.logpdf(model, chain, aggregate = false)

instead of

Turing.logpdf(model, vi, aggregate = false)

.

vollmersj commented 5 years ago

@trappmartin @tlienart

One way to stay within more restritive family usch as hierarchial GLMs

We could use one of them as an example. IMHO rstanarm looks as the easiest to adapt.

trappmartin commented 5 years ago

I just discussed the integration with people at CBL. I’ll compile an issue with some ideas. Making a list of simple models that we support + example implementations sounds like a good plan!

I’ll link the issue here ASAP.

JuliaAI / MLJ.jl

Integrate Turing.jl probablistic programming #157