gaurav-arya / StochasticAD.jl

Research package for automatic differentiation of programs containing discrete randomness.
MIT License
192 stars 15 forks source link

method error with Distributions.BinomialGeomSampler when implementing new particle filter #45

Open slwu89 opened 1 year ago

slwu89 commented 1 year ago

Hi!

I'm trying to modify the particle filter example for a simple SIR model in discrete time (a Markov chain). When using ForwardDiff.gradient, I get the following MethodError:

ERROR: MethodError: no method matching Distributions.BinomialGeomSampler(::Int64, ::ForwardDiff.Dual{ForwardDiff.Tag{var"#106#107", Float64}, Float64, 2})

It arises from the line where the dynamics are applied to each particle to update their trajectory X = map(x -> dyn(x,p,Δt), X). I've posted my working example to a gist here: https://gist.github.com/slwu89/a165175c39cdf6ff9744df892238ac03

If you have any time to take a look or suggest changes I'd really appreciate it. Thanks.

mschauer commented 1 year ago

Our example had continuous dynamics and only the resampling step was discrete. In your example the resampling step is also discrete. The smoothing backend which will take care of this is not yet merged, but we are working on that. In the mean time, maybe it helps converting the Dual into a stochastic triple, sampling the discrete step and smoothing the result back on a dual works. We can look into this together.

slwu89 commented 1 year ago

Thanks @mschauer, yes both the dynamics and resampling are discrete in my case. Ok, thanks for the advice! To be specific, do you mean doing the conversion right before the particle filter calls the dyn function to run the dynamics of the model? And to smooth the result, you recommend to use the stop-gradient method?

slwu89 commented 1 year ago

@mschauer interestingly my SIR example (with the SIR discrete resampling dynamics unchanged) works when taking the derivative of a simple function which returns the total epidemic size (size of the R compartment after the transient dynamics end). Using StochasticAD.derivative_estimate returns a reasonable answer (derivative of final epidemic size wrt contact term is positive and recovery term is negative).

using StochasticAD
using Distributions
using DistributionsAD

@inline function rate_to_proportion(r, t)
    1-exp(-r*t)
end

function dyn(x, p, Δt)
    S,I,R = x
    β,γ = p
    N = S+I+R
    ifrac = rate_to_proportion(β*I/N,Δt)
    rfrac = rate_to_proportion(γ,Δt)
    infection = rand(Binomial(S,ifrac))
    recovery = rand(Binomial(I,rfrac))
    return [S - infection, I + infection - recovery, R + recovery]
end

function total_epidemic_size(nsteps, x0, p, Δt)
    x = copy(x0)
    for n in 2:nsteps
        x = dyn(x, p, Δt)
    end
    return x[3]
end

const p = [0.5, 0.25]
const Δt = 0.1
const x0 = [990, 10, 0]

nsteps = 400
tmax = nsteps*Δt

StochasticAD.derivative_estimate((x) -> total_epidemic_size(nsteps, x0, x, Δt), p)
gaurav-arya commented 1 year ago

Your implementation above, with concurrent discrete random steps for S I and R, should be unproblematic for stochastic triples 🙂

slwu89 commented 1 year ago

Hi @gaurav-arya! The particle filter in the github gist or the small thing above just summing up the total epidemic size? I now understand the issue was because both the model's dynamics were discrete and the resampling step in the particle filter was discrete too. Anyway, I'll rebuild the latest version of StochasticAD and see if it works now!

gaurav-arya commented 1 year ago

Sorry, I was looking at just your small snippet above, not the full particle filter you originally sent. In #70 I've added new_weight support for stochastic triples, so that the unbiased approach can be used with them too.

However, you're indeed correct that discrete randomness in both the model dynamics and the resampling step of a particle filter simultaneously is something that we have not explored in depth with stochastic triples. I did try to get smoothing via ForwardDiff to go through for your example, however smoothing the (integer) number of trials parameter in the Binomial proved tricky to do with duals; something to revisit in the future!

slwu89 commented 1 year ago

@gaurav-arya is this something I could potentially help with? It's been a few months since I read your paper describing the methods so I'll have to reread that, but in the meantime are there any places in the code to check out?

gaurav-arya commented 1 year ago

Hey! Sorry for the late reply here. If you're interested, I think one interesting direction to go in would be to push through the smoothing approach to work on the particle filter. The tricky part would be getting something like this to work:

function f(p)
   y = rand(Binomial(10, p))
   z = rand(Binomial(y, p))
end

ForwardDiff.derivative(f, 0.5) # throws an InexactError

A couple places to look in the code are the ForwardDiff rules and this hack to get stochastic triples to work with the n of a binomial.

I also do hope to get a stochastic triple-based smoothing approach implemented soon, which should be more flexible in terms of propagating through discrete constructs -- once I've done that I could also let you know if you'd like to play around with it for the particle filter:)

slwu89 commented 1 year ago

Thanks @gaurav-arya! Okay, I'll look around in those places you suggested. Also, I'll keep on the lookout for your update, which may happen before I get time to check out the code in detail again.

slwu89 commented 12 months ago

@gaurav-arya I'm somewhat coming back to this, it looks to me as if that "hack" is already implemented for perturbing the n of a binomial? I'm looking at this: https://github.com/gaurav-arya/StochasticAD.jl/blob/main/src/discrete_randomness.jl#L127 Could you describe in a little more detail exactly what needs to happen to get that working?

Also, I somehow accidentally deleted the gist with the original SIR particle filter, I recoded a new one here: https://gist.github.com/slwu89/f6d0c595cea9cc1a86acda470709b140