cscherrer / Soss.jl

Probabilistic programming via source rewriting
https://cscherrer.github.io/Soss.jl/stable/
MIT License
414 stars 30 forks source link

Modeling discrete variables #315

Open jariji opened 2 years ago

jariji commented 2 years ago

Hi @cscherrer,

I tried writing Stan's estimating parameters of a mixture in Soss but I must be missing something. I'm not sure how to translate the loop

for (n in 1:N) {
    vector[K] lps = log_theta;
    for (k in 1:K) {
      lps[k] += normal_lpdf(y[n] | mu[k], sigma[k]);
    }
    target += log_sum_exp(lps);
}
# https://mc-stan.org/docs/2_28/stan-users-guide/summing-out-the-responsibility-parameter.html

using TransformVariables, LogDensityProblems, DynamicHMC,
    DynamicHMC.Diagnostics, Parameters, Statistics, Random, Distributions, Soss

model = @model (N,K) begin
    # Replaced LogNormal with HalfCauchy.
    σ .~ fill(HalfCauchy(2.), K)
    μ .~ fill(Normal(0., 10.), K)
    y ~ For(1:N) do n
        # What about `lps`?
        For(1:K) do k
            Normal(μ[k], σ[k])
        end
        # What about `target?`
    end
end
sourceLogdensity(model)

Manually writing

For(1:N) do n
        lps = log.(θ)
        For(1:K) do k
            lps[k] += logpdf(Normal(μ[k], σ[k]), y[n])
        end
        _ℓ += log(sum(exp.(lps)))
    end

doesn't seem right, since I thought _ℓ gets defined in sourceLogdensity, not in the model.