Posterior prediction of expected value or selected model variables (e.g. mu) (like posterior_epred() and posterior_linpred() in brms)

[JohannesNE]

R's brms package has posterior_epred() , posterior_linpred() and posterior_predict(). These functions sample the posterior for the expected value ( E(y) ), linear predictor (e.g. μ) and entire posterior distribution (y).

See this illustration from Andrew Heiss's blog

The entire posterior is already implemented (see ?predict()), but epred and linpred does not seem to be.

Could this be implemented in predict(), e.g.

using Turing

@model function linear_reg(x, y, σ = 0.1)
    β ~ Normal(0, 1)
    
    # everything with the @linpred macro is exported as columns in an array then predict(; type="linpred")
    @linpred μ = β .* x 

    for i ∈ eachindex(y)
        y[i] ~ Normal(μ[i], σ)
    end
end

f(x) = 2 * x + 0.1 * randn()

Δ = 0.1
xs_train = 0:Δ:10; ys_train = f.(xs_train);
xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

m_lin_reg = linear_reg(xs_train, ys_train);
chain_lin_reg = sample(m_lin_reg, NUTS(100, 0.65), 200);

m_lin_reg_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)));

predictions = predict(m_lin_reg_test, chain_lin_reg, type="linpred") # or a better way to select type

Expected values are probably more difficult as they depend on the likelihood distribution.

[torfjelde]

We can't really implement posterior_linpred for arbitrary models since it assumes a very specific form of likelihood.

posterior_epred is technically possible, but it's a bit unclear to me if it's worth its own method given that it's fairly simple to implement using predict (which is equivalent to your posterior_predict)?

With that being said, you can get posterior_linpred fairly easily using the generated_quantities method:)

Taking your example:

using Turing

@model function linear_reg(x, y, σ = 0.1)
    β ~ Normal(0, 1)
    μ = β .* x 

    for i ∈ eachindex(y)
        y[i] ~ Normal(μ[i], σ)
    end

    # `generated_quantities` simply extracts whatever is present in `return`
    # conditioned on samples in a `chain`.
    return (linpred = μ, pred = y)
end

f(x) = 2 * x + 0.1 * randn()

Δ = 0.1
xs_train = 0:Δ:10; ys_train = f.(xs_train);
xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

m_lin_reg = linear_reg(xs_train, ys_train);
chain_lin_reg = sample(m_lin_reg, NUTS(100, 0.65), 200);

m_lin_reg_test = linear_reg(xs_test, Vector{Union{Missing, Float64}}(undef, length(ys_test)));

# Call `generated_quantities` to get the `linpred` and `pred` from `return`.
generated_quantities(m_lin_reg_test, chain_lin_reg)

[JohannesNE]

Thank you @torfjelde ! generated_quantities is exactly what I was looking for (save for epred, but i can calculate that).
I struggled a bit with the returned array of tuples, but got it working by converting it to a matrix for each quantity.

function get_sample_matrix(gq, selector)
    reduce(vcat, getindex(samp, selector)' for samp in gq[:,1])
end

(also found this, but did not get it working https://gist.github.com/torfjelde/37be5a672d29e473983b8e82b45c2e41)