Help with getting SMC to work on my example?

[sdwfrost]

Hi everyone, thanks for Gen.jl! I've been having fun with it. I've been trying to get my head around the SMC example, and applying it to my problem. Can anyone give any pointers? Code using Metropolis Hastings is below; you can also find a hand-coded particle filter here for reference.

using OrdinaryDiffEq
using SciMLSensitivity
using Distributions
using Gen
using GenDistributions

# Define the Gen model
@gen function bayes_sir_markov(l::Int=40,
                               N::Int=1000,
                               c::Float64=10.0,
                               γ::Float64=0.25,
                               δt::Float64=1.0)
    i₀ ~ uniform_discrete(1, 100)
    β ~ uniform(0.01, 0.1)
    S = N - i₀
    I = i₀
    for i in 1:l
        ifrac = 1-exp(-β*c*I/N*δt)
        rfrac = 1-exp(-γ*δt)
        infection = {(:y, i)} ~ binom(S,ifrac)
        recovery = {(:z, i)} ~ binom(I,rfrac)
        S = S-infection
        I = I+infection-recovery
    end
    nothing
end

# Set the true parameters and solve the model
p = Gen.choicemap()
p[:β] = 0.05
p[:i₀] = 10
l = 40
(sol, _) = Gen.generate(bayes_sir_markov, (l,), p)

# Generate observations
Y = [sol[(:y, i)] for i=1:l]
observations = Gen.choicemap()
for (i, y) in enumerate(Y)
    observations[(:y, i)] = y
end

## Metropolis-Hastings

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


@gen function sir_proposal(current_trace)
    β ~ truncated_normal(current_trace[:β], 0.01, 0.0, Inf)
    i₀ ~ uniform_discrete(current_trace[:i₀]-1, current_trace[:i₀]+1)
end

n_iter = 110000
burnin =  10000
thin = 10
β_mh = Vector{Real}(undef, (n_iter - burnin) ÷ thin)
i₀_mh = Vector{Int}(undef, (n_iter - burnin) ÷ thin)

(tr,) = Gen.generate(bayes_sir_markov, (l,), merge(observations, p))
for i in 1:n_iter
    (tr, _) = Gen.mh(tr, sir_proposal, ())
    if i <= burnin
        continue
    end
    if i % thin == 0
        β_mh[(i - burnin) ÷ thin] = tr[:β]
        i₀_mh[(i - burnin) ÷ thin] = tr[:i₀]
    end
end

# Posterior mean estimates
mean(β_mh) # cf 0.05
mean(i₀_mh) # cf 10

[ztangent]

Do you already have a partial implementation of SMC for this model?

In addition to what's in the docs and the Gen tutorial on particle filtering, you might be interested in using this SMC extension package for Gen called GenParticleFilters.jl -- there's an example in the README for that package that is pretty similar to your hand-coded particle filter.

[sdwfrost]

Thanks for the tips @ztangent! I guess my main question is that for the model implementation above, my model doesn't return anything. In the SMC case, what (if anything) do I have to return?

[ztangent]

Happy to help! And interesting results, thanks! I wonder if the difference in performance is simply because both importance sampling and MCMC are using 10 to 100x more computation than the current SMC algorithm. If you ran 1000 replicates of SMC, each using 1000 particles, and then draw one particle from each replicate using Gen.sample_unweighted_traces, I would expect this to be better than just using importance_resampling.

For variational inference, the current support in Gen.jl is not as strong as we'd like, but there are some examples of black-box variational inference in the test cases: https://github.com/probcomp/Gen.jl/blob/master/test/inference/variational.jl

[sdwfrost]

Just using 100 replicates (due to time), with 1000 particles, I still get quite a few outliers. Can I ask why to use Gen.sample_unweighted_traces, rather than use the weights?

# Iterative SMC

num_replicates = 100
num_particles = 1000
β_smc2 = Vector{Real}(undef, num_replicates)
i₀_smc2 = Vector{Real}(undef, num_replicates)
for i in 1:num_replicates
    state = particle_filter(Y, num_particles)
    trace = sample_unweighted_traces(state,1)[1]
    β_smc2[i] = trace[:β]
    i₀_smc2[i] = trace[:i₀]
end

mean(β_smc2),sqrt(var(β_smc2)) # (0.049592288660781, 0.007121180591284114)
mean(i₀_smc2),sqrt(var(i₀_smc2)) # (14.391547526358067, 17.99118514431295)

histogram(β_smc2)
histogram(i₀_smc2)

[ztangent]

Interesting that you're still getting outliers! I suspect that you are actually getting particle filter collapse a non-trivial amount of time, leading to the outliers. Effective sample size isn't enough to track whether that has happened, since after you perform a resampling move, the weights of the particles all go back to being equal, which means effective sample size goes back to the number of particles in the filter --- even if all of them are identical copies of each other. So what might be happening here is that on least some runs, the particle filter collapses on a value of i₀ that initially seems good after e.g. 5 observations, but turns out to be a bad explanation of the remaining 35 observations. Since the particle filter currently doesn't have any rejuvenation moves that perturb the value of i₀, there's no way to recover from this situation.

A couple of ways to address would be to:

• Remove the resampling step altogether (this makes the particle filter equivalent to important sampling, except with slower runtime if you don't use the Unfold combinator)
• Resample less regularly --- e.g. check the resampling condition only every 5 or 10 timesteps, instead of every timestep.
• Add MCMC rejuvenation moves every few timesteps.

As for why to use Gen.sample_unweighted_traces, that's just to make a fair comparison with Gen.importance_resampling. Gen.importance_resampling first samples a collection of 1000 importance-weighted traces, then finally samples one trace from that collection according to its normalized importance weight. The particle filter also produces such a collection of 1000 importance-weighted traces. You can either use that weighted collection as your approximation to the posterior --- or, you can draw one particle from it (according to its normalized particle weight), giving you an unweighted sample that should look like it comes from something close to the true posterior. Doing this multiple times mimics what you've done with importance resampling.

[sdwfrost]

Here is my example to date - I want to show packages in their best light, so suggestions to improve are much appreciated!