Turing.jl Gibbs with closed-form conditionals

The new version for the Gibbs sampler is missing a feature that was present in the previous version. It would be nice to implement an equivalent (or improved!) feature.

Interface

GibbsConditional(sym, conditional)

A "pseudo-sampler" to manually provide analytical Gibbs conditionals to Gibbs. GibbsConditional(:x, cond) will sample the variable x according to the conditional cond, which must therefore be a function from a NamedTuple of the conditioned variables to a Distribution.

Examples

α_0 = 2.0
θ_0 = inv(3.0)
x = [1.5, 2.0]
N = length(x)

@model function inverse_gdemo(x)
    λ ~ Gamma(α_0, θ_0)
    σ = sqrt(1 / λ)
    m ~ Normal(0, σ)
    @. x ~ \$(Normal(m, σ))
end

# The conditionals can be formulated in terms of the following statistics:
x_bar = mean(x) # sample mean
s2 = var(x; mean=x_bar, corrected=false) # sample variance
m_n = N * x_bar / (N + 1)

function cond_m(c)
    λ_n = c.λ * (N + 1)
    σ_n = sqrt(1 / λ_n)
    return Normal(m_n, σ_n)
end

function cond_λ(c)
    α_n = α_0 + (N - 1) / 2 + 1
    β_n = s2 * N / 2 + c.m^2 / 2 + inv(θ_0)
    return Gamma(α_n, inv(β_n))
end

m = inverse_gdemo(x)

sample(m, Gibbs(GibbsConditional(:λ, cond_λ), GibbsConditional(:m, cond_m)), 10)