Turing.jl programmatically set number of clusters in a mixture model

Specific Domains Probabilistic programming
1

I am trying to model some data using a mixture model with Turing.jl, but I want to programmatically set the number of Gaussian in the mixture without having to edit the model. For example if I want ten Gaussians I don’t want to manually add 10
mu ~ Normal() statements.

Currently my working ‘manual’ model looks like:

using Turing, Distributions

@model function gmm()
   try
       ω ~ Dirichlet(3, 1.0)


       μ1 ~ MvNormal(zeros(2), I)
       Σ1 ~ InverseWishart(2, [1 0;0 1])
  
       μ2 ~ MvNormal(zeros(2), I)
       Σ2 ~ InverseWishart(2, [1 0;0 1])


       data ~ MixtureModel(
           [
               Product(fill(Uniform(-10,10), 2)),
               MvNormal(μ1, Σ1),
               MvNormal(μ2, Σ2)
           ],
           ω
       )
   catch me
       if me isa PosDefException
           Turing.@addlogprob! -Inf
       end
   end
end

I saw others recommend using something like the following:

using Turing, LinearAlgebra, Distributions

@model function test_mixture_model(;K = 3)

    try
        ω ~ Dirichlet(K, 1.0)

        μ ~ filldist(MvNormal(5 .*ones(2), 5*I), K-1)
    
        obs ~ MixtureModel(
            [
                id == 1 ? Product(fill(Uniform(-10, 10), 2)) : MvNormal(μ[id-1], I) for id in 1:K
            ],
            ω
        )

    catch me 
        if me isa PosDefException
            Turing.@addlogprob! -Inf
        end
    end
end

model = test_mixture_model()

ϕ = MixtureModel([Product(fill(Uniform(-10,10), 2)), MvNormal(zeros(2), I), MvNormal([7,7], I)])
data_synth = rand(ϕ,1_000)
model_cond = model | (; obs = data_synth)
chain = sample(model_cond, NUTS(; adtype=AutoForwardDiff()), 1000)

However, I encountered an issue with all the μs becoming tied together, i.e. they contain similar values even though the ground truth has the clusters far apart. Then If I try to add Σ as a parameter they don’t show up as inferred values.

Summary Statistics
  parameters      mean       std      mcse    ess_bulk   ess_tail      rhat   ess_per_sec 
      Symbol   Float64   Float64   Float64     Float64    Float64   Float64       Float64 

        ω[1]    0.3448    0.2338    0.0073   1037.2750   695.0540    1.0017     2160.9896
        ω[2]    0.3301    0.2335    0.0078    864.5741   679.1482    1.0019     1801.1959
        ω[3]    0.3251    0.2315    0.0077    845.4391   809.9245    1.0006     1761.3315
     μ[1, 1]    4.9690    2.0700    0.0555   1356.5816   901.3844    1.0024     2826.2116
     μ[2, 1]    4.9583    2.1976    0.0582   1411.3363   785.8980    1.0046     2940.2841
     μ[1, 2]    5.0194    2.2868    0.0686   1111.6677   861.5815    0.9997     2315.9744
     μ[2, 2]    5.0070    2.1828    0.0701    967.7597   786.9313    0.9993     2016.1660

To summarize, I would like to have a mixture model where the user can select the number of Gaussian to include. Additionally I would like the mean and covariance of those Gaussian to be inferred by the turing from the data. Is there a recommend way of doing this?