Support for Stan, Soss, and Gen

[ParadaCarleton]

Currently, only Turing is supported; support for Stan, Soss, and Gen would be great.

[ParadaCarleton]

@itsdfish @goedman if either of you happen to have any time and want to add Stan support, let me know; I'd like to make the package announcement as soon as I have Stan support.

[goedman]

@ParadaCarleton (and @itsdfish) Carlos, I am wondering what you mean by this. I'm using it all the time now with Stan and it is very simple to generate the required input data for ParetoSmooth.jl.

Unfortunately, in Stan you typically generate the required ParetoSmooth input data (I usually call it the log_lik 3d matrix) by adding a generated quantities section at the end of the model, which is different from Turing's way.

Yes, for StatisticalRethinking I need to do a bit extra and will complete that work in v4 of StatisticalRethinking (and that is primarily an issue of presenting the results).

It is possible to add an example to the tests but we discussed that a while ago and it would require building Stan's cmdstan executable when running the tests on CI which I think is too much of a hassle.

It is also possible to add a doc section with an example of how to use ParetoSmooth from StatisticalRethinkingJulia.

[itsdfish]

@goedman, perhaps @ParadaCarleton wants something similar to R. It looks like R uses some sort of model object returned from Stan to compute loo. My guess is that is uses expose_stan_functions to expose the C++ likelihood function to R. One hack might be to use RCall, but that could prove to be fragile. Something like this is outside of my expertise.

[goedman]

@itsdfish ( @ParadaCarleton )

Thanks Chris. Currently theloo_compare() method I added to StatisticalRethinking.jl is:

function loo_compare(models::Vector{SampleModel}; 
    loglikelihood_name="log_lik",
    model_names=nothing,
    sort_models=true, 
    show_psis=true)

    if isnothing(model_names)
        mnames = [models[i].name for i in 1:length(models)]
    end

    nmodels = length(models)

    ka = Vector{KeyedArray}(undef, nmodels)
    ll = Vector{Array{Float64, 3}}(undef, nmodels)
    psis = Vector{PsisLoo{Float64, Array{Float64, 3},
        Vector{Float64}, Int64, Vector{Int64}}}(undef, nmodels)

    for i in 1:length(models)
        ka[i] = read_samples(models[i], :keyedarray)
        ll[i] = permutedims(Array(matrix(ka[i], loglikelihood_name)), [3, 1, 2])
        psis[i] = psis_loo(ll[i])
        show_psis && psis[i] |> display
    end

    loo_compare(psis...; model_names=mnames, sort_models)
end

So in a call such as loo_compare([m5_1s, m5_2s, m5_3s]) I pass in very similar (equivalent?) objects as in R and prepare it to call loo_compare in ParetoSmooth.

It's very likely I can (and will) improve above formulation as I do further work on StatisticalRethinking v4. But that will be hidden from users. It's also tempting to add a loo_compare method passing in a vector of KeyedArray chain objects as in many cases that will be returned from previous calls to read_samples().

I do know Richard McElreath adds a generated quantities block to the Stan model when the log_lik matrix is needed for model comparison. In the latest Stan docs (for the new cmdstanR option) a listed goal is to prevent directly calling C++, just executables. And as you know in StanJulia I have always stayed away from the expose_stan_functions approach.

Just my $0.02.

[itsdfish]

@goedman, I understand your reluctance to work with the C++ code. Even in the example with R, they use generated quantities as you noted. So there is a minor step of extracting the log likelihoods from the Stan output. That leads me to believe dealing with the C++ code would be more trouble than it is worth.

Would you mind providing a link to code where you extract the log likelihoods from a Stan model and pass it to loo in Paretosmooth? That might allow us to determine whether there is any room for a convenience function. If not, we can use your example in the documentation as an example with Stan.

[goedman]

Hi @itsdfish ,

In above method I do (for each model):

# By default read_samples() returns a KeyedArray based chain
chns = read_samples(m5_1s)

# Extract loglik KeyedArray and permute dimensions (what you call "extract the log likelihoods from a Stan model")
# `:log_lik` is the default and can be dropped
# In StanSample.jl matrix() is overloaded from the Tables.jl API (supported by KeyedArrays)
loglik = permutedims(Array(matrix(chns, :log_lik)), [3, 1, 2])

# Compute a PsisLoo object
psisloo = psis_loo(loglik)

# A vector of PsisLoo objects is passed to loo_compare() in ParetoSmooth.

Above method is used here.

[ParadaCarleton]

It is possible to add an example to the tests but we discussed that a while ago and it would require building Stan's cmdstan executable when running the tests on CI which I think is too much of a hassle.

Oh, thanks -- I didn't notice that. Makes sense. That being said, doesn't the StatisticalRethinking library use Stan? I believe it's possible to check whether your package breaks another package's tests using CI. @yebai mentioned something about testing whether changes in AdvancedHMC break Turing. Or perhaps Turing already runs tests comparing Turing and Stan, in which case we could maybe piggyback off of that?

[ParadaCarleton]

@itsdfish ( @ParadaCarleton )

Thanks Chris. Currently theloo_compare() method I added to StatisticalRethinking.jl is:

function loo_compare(models::Vector{SampleModel}; 
    loglikelihood_name="log_lik",
    model_names=nothing,
    sort_models=true, 
    show_psis=true)

    if isnothing(model_names)
        mnames = [models[i].name for i in 1:length(models)]
    end

    nmodels = length(models)

    ka = Vector{KeyedArray}(undef, nmodels)
    ll = Vector{Array{Float64, 3}}(undef, nmodels)
    psis = Vector{PsisLoo{Float64, Array{Float64, 3},
        Vector{Float64}, Int64, Vector{Int64}}}(undef, nmodels)

    for i in 1:length(models)
        ka[i] = read_samples(models[i], :keyedarray)
        ll[i] = permutedims(Array(matrix(ka[i], loglikelihood_name)), [3, 1, 2])
        psis[i] = psis_loo(ll[i])
        show_psis && psis[i] |> display
    end

    loo_compare(psis...; model_names=mnames, sort_models)
end

So in a call such as loo_compare([m5_1s, m5_2s, m5_3s]) I pass in very similar (equivalent?) objects as in R and prepare it to call loo_compare in ParetoSmooth.

It's very likely I can (and will) improve above formulation as I do further work on StatisticalRethinking v4. But that will be hidden from users. It's also tempting to add a loo_compare method passing in a vector of KeyedArray chain objects as in many cases that will be returned from previous calls to read_samples().

I do know Richard McElreath adds a generated quantities block to the Stan model when the log_lik matrix is needed for model comparison. In the latest Stan docs (for the new cmdstanR option) a listed goal is to prevent directly calling C++, just executables. And as you know in StanJulia I have always stayed away from the expose_stan_functions approach.

Just my $0.02.

Something like this would actually be pretty much exactly what I need! I think with a bit of cleanup this would work -- we'd just need to fix the overly-strict typing. Apart from that I think we'd also need a method to perform LOO-CV on a single model, but that's about it.

[goedman]

@itsdfish & @ParadaCarleton

This is likely how it is going to end up in StatisticalRethinking v4 (provided StanSample v4 has been loaded):

# KeyedArray chains are [draws, chains, params] while ParetoSmooth
# expects [params, draws, chains].
function to_paretosmooth(ll, pd = [3, 1, 2])
    permutedims(ll, pd)
end

function loo_compare(models::Vector{SampleModel}; 
    loglikelihood_name="log_lik",
    model_names=nothing,
    sort_models=true, 
    show_psis=true)

    nmodels = length(models)
    model_names = [models[i].name for i in 1:nmodels]

    chains_vec = read_samples.(models) # Obtain KeyedArray chains
    loo_compare(chains_vec; loglikelihood_name, model_names, sort_models, show_psis)
end

function loo_compare(chains_vec::Vector{<: KeyedArray}; 
    loglikelihood_name="log_lik",
    model_names=nothing,
    sort_models=true, 
    show_psis=true)

    nmodels = length(chains_vec)

    ll_vec = Array.(matrix.(chains_vec, loglikelihood_name)) # Extract log_lik matrix
    ll_vecp = map(to_paretosmooth, ll_vec) # Permute dims for ParetoSmooth
    psis_vec = psis_loo.(ll_vecp) # Compute PsisLoo for all models

    if show_psis # If a printout is needed
        for i in 1:nmodels
            psis_vec[i] |> display
        end
    end

    # Call vanilla loo_compare() in ParetoSmooth.jl
    loo_compare(psis_vec...; model_names, sort_models)
end

export
    to_paretosmooth,
    loo_compare

I haven't worked out the details for StatisticalRethinking v4 when Turing is loaded. But that will mostly follow what Chris did.

[goedman]

In StatisticalRethinking v4, with StanSample loaded, for LOO_CV I will use:

function psis_loo(model::SampleModel; loglikelihood_name="log_lik")
    chains = read_samples(model) # Obtain KeyedArray chains
    psis_loo(chains; loglikelihood_name)
end

function psis_loo(chains::T; loglikelihood_name="log_lik") where {T <: KeyedArray}
    ll = Array(matrix(chains, loglikelihood_name)) # Extract log_lik matrix
    ll_p = to_paretosmooth(ll) # Permute dims for ParetoSmooth
    psis = psis_loo(ll_p) # Compute PsisLoo for model
end

which typically is called as:

chns5_1s = read_samples(m5_1s)
psis_loo(chns5_1s) |> display

[ParadaCarleton]

@goedman I just had a thought -- would you be willing to add that code to StanSample.jl with a couple tests, then? (And Requires.jl so that you don't have the dependency on ParetoSmooth.jl, I assume.) That way we can have out-of-the-box model comparison functionality in Stan.

[goedman]

Hi Carlos (@ParadaCarleton),

Apologies for the delay in answering.

StanSample is a relative lightweight package on a lower level (it has no notion of comparing models). I would prefer to keep it that way.

ParetoSmooth based model comparison really fits into (the very much simplified) StatisticalRethinking.jl v4.0. The simplifications come from dropping usage of @reexports ... and all graphics

The code is included here and tested here in SR v4.0.3.

As I mentioned before, for the SR course clips in StatisticalRethinkingStan I will follow the book and present the results in a different format and re-introduce some columns and fields into the SR's version ModelComparison object (likely more like the old LooCompare object), also to improve presentation in Pluto notebooks (e.g. the Pareto_k value warnings currently show up in the console).

All graphics are now handled in StatisticalRethinkingPlots which does contain an option to plot the Pareto_k values. And in the future similarly in StatisticalRethinkingMakie.