Help with deterministic traces

[sdwfrost]

I'm trying to simulate/fit a problem where I have access to a numerical CDF, and thought I could use the inverse CDF approach to sample from this distribution in Gen.jl. As a simple test case, let's try to fit a normal distribution, firstly by drawing samples from a normal, and then by drawing uniform samples and using the inverse CDF. I've looked up the documentation on deterministic traces, but I'm unclear where to go next such that I could use Gen.mh etc. What are the next steps below? I tried to comment my code throughout.

using Gen
using Distributions
using GenDistributions

@gen function normal_test(K::Int64)
    # Sample from normal
    μ ~ normal(0.0, 1.0)
    σ ~ gamma(1.0, 1.0)
    for i in 1:K
        {(:xi, i)} ~ normal(μ, σ)
    end
end

# Set parameters for observed data
p = Gen.choicemap()
p[:μ] = 0.0
p[:σ] = 1.0
K = 10000
# Simulate 'true' data
(sol, _) = Gen.generate(normal_test, (K,), p)
x = [sol[(:xi,i)] for i in 1:K]
# Fix observations
xobs = choicemap()
for (i, xi) in enumerate(x)
    xobs[(:xi, i)] = xi
end

# Define truncated normal proposal
const truncated_normal = DistributionsBacked((mu,std,lb,ub) -> Distributions.Truncated(Normal(mu, std), lb, ub),
                                             (true,true,false,false),
                                             true,
                                             Float64)

# Define a proposal function
@gen function normal_proposal(current_trace)
    μ ~ normal(current_trace[:μ], 0.01)
    σ ~ truncated_normal(current_trace[:σ], 0.01, 0.0, Inf)
end

# MH loop
n_accept = 0
n_iter = 11000
μ_mh = Vector{Real}(undef, n_iter)
σ_mh = Vector{Real}(undef, n_iter)
scores = Vector{Float64}(undef, n_iter)
(tr,) = Gen.generate(normal_test, (K,), merge(xobs, p))
@time for i in 1:n_iter
    global (tr, did_accept) = Gen.mh(tr, normal_proposal, ())
    μ_mh[i] = tr[:μ]
    σ_mh[i] = tr[:σ]
    scores[i] = Gen.get_score(tr)
    if did_accept
        global n_accept += 1
    end
    if i % 100 == 0
        @info "Iteration $i: μ=$(μ_mh[i]), σ=$(σ_mh[i]), score=$(scores[i]), acceptance rate=$(n_accept/i)"
    end
end

# Compute summary statistics
mean(μ_mh), std(μ_mh), mean(σ_mh), std(σ_mh)

# Now try deterministic trace
# [0,1] to normal
@transform f(K) (t1) to (t2) begin
    for i in 1:K 
        xit = @read(t1[(:xit, i)], :continuous)
        xi = quantile(Normal(t1[:μ], t1[:σ]), xit)
        @write(t2[(:xi, i)], xi, :continuous)
    end
end

# Normal to [0,1]
@transform finv(K) (t2) to (t1) begin
    for i in 1:K 
        xi = @read(t2[(:xi, i)], :continuous)
        xit = cdf(Normal(t2[:μ], t2[:σ]), xi)
        @write(t1[(:xit, i)], xit, :continuous)
    end
end

# Specify bijections
pair_bijections!(f, finv)

# Now rewrite model
@gen function normal_test2(K::Int64)
    # Sample from normal
    μ ~ normal(0.0, 1.0)
    σ ~ gamma(1.0, 1.0)
    for i in 1:K
        {(:xit, i)} ~ uniform(0.0, 1.0)
        # What next??
    end
end

Apart from not knowing how to include a transformed trace inside a @gen model, it isn't clear whether the parameters for the transform should come from the trace or passed as parameters when included inside the @gen macro.

[ztangent]

Finally getting to this --- I don't think this is the intended use case for deterministic trace translators, since it's going to add a whole bunch overhead compared to just implementing the inverse CDF sampling trick in Julia directly. The intended usage for (deterministic) trace translators is to perform SMC between models with different structures (see the paper that introduced the trace translator concept: https://files.sri.inf.ethz.ch/website/papers/pldi18_interemental_inference_for_probabilistic_programs.pdf).