fauconab / PolygenicRiskScores.jl

Polygenic Risk Scoring in pure Julia
Other
10 stars 2 forks source link

Turing implementation #14

Open cpfiffer opened 2 months ago

cpfiffer commented 2 months ago

Here's a rough sketch of a Turing version of the model, so you don't have to use a hand-rolled Gibbs sampler:

using Turing

@model function PRSCSxFUCK(ys, Xs, M, j_inds; a=1.0, b=0.5)
    # Size extraction
    k = length(ys)
    n_ks = map(length(ys))
    m_ks = map(x -> size(x, 2), Xs)

    # Sigma
    σ ~ filldist(InverseGamma(2, 3), k)

    # Draw 
    δ ~ filldist(Gamma(b, ϕ), M)
    ψ ~ arraydist([Gamma(a, d) for d in δ], M)

    # 
    βk = Vector{Vector{Float64}}(undef, k)
    for i in 1:k
        # associated betas
        assc_psi = ψ[j_inds[i]]
        var_term = assc_psi .* σ[i]^2 ./ n_ks[i]

        βk[i] ~ MvNormal(zeros(m_ks[i]), diagm(var_term))

        # Likelihood
        ys[i] ~ MvNormal(Xs[i] * βk[i], σ[i])
    end

end

Deeply untested, more of a sketch than anything. There's also a type instability issue with this line

βk = Vector{Vector{Float64}}(undef, k)

but that can be fixed later.