Help with importance resampling

[sdwfrost]

Hi all,

I have a model implementated in Gen.jl, where Metropolis Hastings works well (giving a reasonable posterior), but where importance sampling does very poorly in terms of the posterior distribution. In a not-too-different model (here), IS and MH have fairly similar posterior distributions. Can anyone see why importance (re)sampling isn't working here?

using OrdinaryDiffEq
using DiffEqCallbacks
using Distributions
using Gen
using GenDistributions
using Random
using Plots

# Utility function to turn an exponential rate into a 1-step proportion

@inline function rate_to_proportion(r)
    1-exp(-r)
end

# Define dynamical model

function sir_map!(du,u,p,t)
    (S,I,C,Y) = u # Susceptible, Infected, Cumulative infections, New infections
    (β,γ) = p
    infection = rate_to_proportion(β*I)*S
    recovery = rate_to_proportion(γ)*I
    @inbounds begin
        du[1] = S-infection
        du[2] = I+infection-recovery
        du[3] = C+infection
        du[4] = infection
    end
    nothing
end

# Wrapper function for the above

function simulate(β, γ, ρ, nsteps)
    tspan = (0, nsteps)
    u0 = [1-ρ, ρ, 0.0, 0.0] # Initial conditions: S, I, C, Y
    p = [β, γ]
    prob = DiscreteProblem(sir_map!, u0, tspan, p)
    sol = solve(prob, FunctionMap())
    τ = sol[end][3] # Final size
    f = sol[4,2:end] # New infections per time step
    f = clamp.(f,0.0,1.0) # Ensure stays in range
    f = f/τ # Normalize frequencies to sum to 1
    return sol, τ, f
end

# Parameter values

nsteps = 40   # Number of time steps to simulate
tspan = (0, nsteps)   # Time span for the simulation
u0 = [0.99, 0.01, 0.0, 0.0] # Initial conditions: S, I, C, Y
β = 0.5 # Infection rate
γ = 0.25 # Recovery rate
ρ = 0.01 # Proportion of individuals initially infected
N = 99 # Number of susceptibles sampled at t=0
M = 1  # Number of initially infected individuals

@gen function sds_discrete_gen(N::Int,     # initial susceptibles followed
                               M::Int,     # initial infected followed
                               nsteps::Int) # number of time steps
    # Priors
    β  ~ uniform(0.0, 1.0)
    γ  ~ uniform(0.0, 1.0)
    ρ  ~ uniform(0.0, 0.1)

    # Simulate
    (sol, τ, f) = simulate(β, γ, ρ, nsteps)     # f : infection-time PMF

    # Sample number of infections
    K  ~ binom(N, τ)                       # total infections

    # Sample infection times
    for i in 1:K
        {(:ti, i)} ~ categorical(f)
        # or {(:ti, i)} ~ uniform_discrete(1, nsteps)  # infection times
    end

    # Sample recovery times
    pγ = rate_to_proportion(γ)
    for i in 1:(M + K)
        {(:delta, i)} ~ geometric(pγ)  # recovery intervals
    end
    return sol
end

# Choice map for parameters

p = Gen.choicemap()
p[:β] = β
p[:γ] = γ
p[:ρ] = ρ

# Simulate 'true' data

(sol, _) = Gen.generate(sds_discrete_gen, (N, M, nsteps), p)

# Make observations

function make_obs(ti_obs::Vector{Int}, delta_obs::Vector{Int})
    obs = choicemap()
    obs[:K] = length(ti_obs)
    for (i, t) in enumerate(ti_obs)
        obs[Symbol(:ti, i)] = t
    end
    for (i, d) in enumerate(delta_obs)
        obs[Symbol(:delta, i)] = d
    end
    return obs
end

ti = [sol[(:ti,i)] for i in 1:sol[:K]]
delta = [sol[(:delta,i)] for i in 1:(M+sol[:K])]
observations = make_obs(ti, delta);

## Metropolis Hastings

### Proposal distribution

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.006, 0.0, Inf)
    γ ~ truncated_normal(current_trace[:γ], 0.006, 0.0, 1.0)
    ρ ~ truncated_normal(current_trace[:ρ], 0.006, 0.0, 1.0)
end

n_iter = 11000
β_mh = Vector{Real}(undef, n_iter)
γ_mh = Vector{Real}(undef, n_iter)
ρ_mh = Vector{Real}(undef, n_iter)
scores = Vector{Float64}(undef, n_iter)
# Set up trace for the loop to write to
(tr,) = Gen.generate(sds_discrete_gen, (N,M,nsteps), merge(observations, p))
n_accept = 0
@time for i in 1:n_iter
    global (tr, did_accept) = Gen.mh(tr, sir_proposal, ())
    β_mh[i] = tr[:β]
    γ_mh[i] = tr[:γ]
    ρ_mh[i] = tr[:ρ]
    scores[i] = Gen.get_score(tr)
    if did_accept
        global n_accept += 1
    end
    if i % 1000 == 0
        @info "Iteration $i: β=$(β_mh[i]), γ=$(γ_mh[i]), ρ=$(ρ_mh[i]), score=$(scores[i]), acceptance rate=$(n_accept/i)"
    end
end

burnin = 1000
pl_β_mh = histogram(β_mh[burnin+1:end], label=false, title="β", ylabel="Density", density=true, xrotation=45)
vline!([sol[:β]], label="True value")
pl_γ_mh = histogram(γ_mh[burnin+1:end], label=false, title="γ", ylabel="Density", density=true, xrotation=45)
vline!([sol[:γ]], label="True value")
pl_ρ_mh = histogram(ρ_mh[burnin+1:end], label=false, title="ρ", ylabel="Density", density=true, xrotation=45)
vline!([sol[:ρ]], label="True value")
plot(pl_β_mh, pl_γ_mh, pl_ρ_mh, layout=(1,3), plot_title="MH sampling")

## Importance sampling

num_particles = 100
num_replicates = 10000
β_is = Vector{Real}(undef, num_replicates)
γ_is = Vector{Real}(undef, num_replicates)
ρ_is = Vector{Real}(undef, num_replicates)
for i in 1:num_replicates
    (trace, lml_est) = Gen.importance_resampling(sds_discrete_gen, (N, M, nsteps), observations, num_particles)
    β_is[i] = trace[:β]
    γ_is[i] = trace[:γ]
    ρ_is[i] = trace[:ρ]
    if i % 100 == 0
        @info "Iteration $i: β=$(β_is[i]), γ=$(γ_is[i]), ρ=$(ρ_is[i]), lml=$(lml_est)"
    end
end

## Plots show a poor fit

pl_β_is = histogram(β_is, label=false, title="β", ylabel="Density", density=true, xrotation=45)
vline!([sol[:β]], label="True value")
pl_γ_is = histogram(γ_is, label=false, title="γ", ylabel="Density", density=true, xrotation=45)
vline!([sol[:γ]], label="True value")
pl_ρ_is = histogram(ρ_is, label=false, title="ρ", ylabel="Density", density=true, xrotation=45)
vline!([sol[:ρ]], label="True value")
plot(pl_β_is, pl_γ_is, pl_ρ_is, layout=(1,3), plot_title="Importance resampling")

The issue doesn't seem to be a lack of proposal. I also thought it may be due to not keeping a copy of the trace; however, the below code also doesn't give good performance:

num_particles = 1000
num_replicates = 1000
β_is = Vector{Real}(undef, num_replicates)
γ_is = Vector{Real}(undef, num_replicates)
ρ_is = Vector{Real}(undef, num_replicates)
(trace,) = Gen.generate(sds_discrete_gen, (N,M,nsteps), merge(observations, p))
for i in 1:num_replicates
    global (trace, lml_est) = Gen.importance_resampling(sds_discrete_gen, (N, M, nsteps), observations, sir_proposal, (trace,), num_particles)
    β_is[i] = trace[:β]
    γ_is[i] = trace[:γ]
    ρ_is[i] = trace[:ρ]
    if i % 100 == 0
        @info "Iteration $i: β=$(β_is[i]), γ=$(γ_is[i]), ρ=$(ρ_is[i]), lml=$(lml_est)"
    end
end

[ztangent]

I think there's an issue with your make_obs function. Currently it's:

function make_obs(ti_obs::Vector{Int}, delta_obs::Vector{Int})
    obs = choicemap()
    obs[:K] = length(ti_obs)
    for (i, t) in enumerate(ti_obs)
        obs[Symbol(:ti, i)] = t
    end
    for (i, d) in enumerate(delta_obs)
        obs[Symbol(:delta, i)] = d
    end
    return obs
end

But the Symbol isn't necessary --- based on how you wrote your model, the address should be a Tuple, not a Symbol. So make_obs should be:

function make_obs(ti_obs::Vector{Int}, delta_obs::Vector{Int})
    obs = choicemap()
    obs[:K] = length(ti_obs)
    for (i, t) in enumerate(ti_obs)
        obs[(:ti, i)] = t
    end
    for (i, d) in enumerate(delta_obs)
        obs[(:delta, i)] = d
    end
    return obs
end

Hopefully that's the only issue (i.e. you weren't conditioning on the right observations -- Gen doesn't always throw an error when you condition on addresses that don't exist in the model). Otherwise the model is a bit too complicated for me to debug without running it.

[sdwfrost]

Thanks! That seems to have mostly fixed the issue - I'm surprised that the MH worked though with observations not correctly specified.