TuringLang / docs

Documentation and tutorials for the Turing language
https://turinglang.org/docs/tutorials/docs-00-getting-started/
MIT License
225 stars 97 forks source link

Tutorial 'Bayesian Mixed Models': PG is ok, Gibbs runs into trouble #436

Open CMoebus opened 3 months ago

CMoebus commented 3 months ago

Hi, As a kind of exercise I try to squeeze the code a little bit. Sampling only with PG(...) seems to give reasonable results. But when I try to use the Gibbs(...) sampler, as the original does it generates an error message rather mysterious to me. I provide the code in a Pluto cell below. It would help a lot if I could get a hint what is going wrong. Thanks a lot, Claus

let N = 60 K = 2 D = 2 w = [ 0.5, 0.5] # w is fixed μ0 = [-3.5 +0.5; # μ0 = [μ0[1], μ0[2]] for data generation +3.5 -0.5] μ1 = [ 0.0 0.0; # μ1 = [μ1[1], μ1[2]] as priors 0.0 0.0]
μ = Array{Float64}(undef, (D, K)) x = Array{Float64}(undef, (D, N))

===============================================================================

# Data Generation
#--------------------------------------------------------------------------------
mixtureModel0 = MixtureModel([MvNormal(μ0[:,k], I) for k in 1:K], w)
#                                           ^==== explicit column vectors 
x = rand(mixtureModel0, N)
#                  ^======================== continuous cluster membership
# ===============================================================================
# Mixture Model
#--------------------------------------------------------------------------------
@model function bayesianMixtureModel(x)
    k = Vector{Int64}(undef, N)
    #----------------------------------------------------------------------------
    for k in 1:K
        # sampling priors μ1 = [μ1[1], μ1[2]]
        μ[:, k] ~ MvNormal(μ1[:,k], 2.0 .* I)
    end # for k
    #
    mixtureModel = [MvNormal(μ[:,k], I) for k in 1:K] # <=== discrete membership
    #                            ^==== explicit column vectors
    for i in 1:N
        k[i]   ~ Categorical(w)                   # k[i], vector is important
        # ^======================== discrete cluster membership
        x[:,i] ~ mixtureModel[k[i]]               # likelihood
    end # for i
end # function bayesianMixtureModel
#--------------------------------------------------------------------------------
model = bayesianMixtureModel(x)
#--------------------------------------------------------------------------------
chains = 
    let nSamples = 100
        # sampler  = Prior()
        # sampler = PG(100, :k, :μ)                         # <========= seems ok 
        sampler = Gibbs(PG(100, :k), HMC(0.05, 10, :μ))
        sample(model, sampler, nSamples)
    end # let
#--------------------------------------------------------------------------------
plot(chains[["μ[:,1][1]", "μ[:,1][2]", "μ[:,2][1]", "μ[:,2][2]"]], colordim = :parameter, legend=true)
#--------------------------------------------------------------------------------
assignments = [mean(chains, "k[$i]") for i in 1:N]
#--------------------------------------------------------------------------------
μMean1 = [mean(chains, "μ[:,1][$i]") for i in 1:D]
μMean2 = [mean(chains, "μ[:,2][$i]") for i in 1:D]
# ===============================================================================
scatter(x[1, :], x[2, :]; label="data point", title=L"Data 3 (Discrete Post. Membership $k[i]$ with Turing)", zcolor=assignments)
#--------------------------------------------------------------------
# plot of the prior cluster centroids
scatter!([(μ0[1, 1], μ0[2, 1]), (μ0[1, 2], μ0[2, 2])], color=:turquoise, label="prior cluster centroid", markersize=6) 
#
# plot of the posterior cluster centroids
scatter!([(μMean1[1], μMean1[2]), (μMean2[1], μMean2[2])], color=:violet, label="posterior cluster centroid", markersize=6)
#--------------------------------------------------------------------------------

end # let