Ibn_paper.jl

using Turing

@model function ibp(y, α, kmax, ::Type{MV}=Vector{Float64}) where {MV}
   N = length(y)
   
   ks = tzeros(Int, N)
   ks[1] ~ Poisson(α)
   ks[1] = ks[1] <= kmax ? ks[1] : kmax
   
   z = tzeros(Int, N, kmax)
   z[1,1:ks[1]] .= 1

   for i in 2:N
       K = sum(ks[1:i-1])
       for j in 1:K
           mk = sum(z[:,j])
           z[i,j] ~ Bernoulli(mk / i)
       end
       ks[i] ~ Poisson(α / i)
       ks[i] = K+ks[i] <= kmax ? ks[i] : 0
       if ks[i] > 0
           z[i,(K+1):sum(ks[1:i])] .= 1
       end
   end

   K = sum(ks)
   μ = MV(undef, K)
   for j = 1:K
       μ[j] ~ Normal(0.0, 10.0)
   end

   for i in 1:N
       y[i] ~ Normal(μ' * z[i,1:K], 1.0)
   end
end

x = [-1.48, -1.40, -1.16, -1.08, -1.02, 0.14, 0.51, 0.53, 0.78];
model = ibp(x, 10.0, 100)

@time sample(model, Gibbs(PG(4,:z,:ks), HMC(0.01, 5, :μ)), 5000, specialize_after=0, chain_type=Any);
# 424.308852 seconds (2.06 G allocations: 97.240 GiB, 3.94% gc time)
@time sample(model, Gibbs(PG(4,:z,:ks), HMC(0.01, 5, :μ)), 5000, specialize_after=1, chain_type=Any);
# 326.290391 seconds (1.57 G allocations: 61.866 GiB, 3.31% gc time)
@time sample(model, Gibbs(PG(4,:z,:ks), HMC(0.01, 5, :μ)), 5000, specialize_after=10, chain_type=Any);
# 367.878543 seconds (1.73 G allocations: 67.285 GiB, 3.40% gc time)