Roadmap on stochastic models?

[sdwfrost]

I saw that stochastic models were on the roadmap, and was curious as to what was meant by that. Will it be possible to fit models such as this one, for which I used Gen.jl?

[ismailsenoz]

By stochastic models, we mean solving stochastic differential/difference equations using RxInfer. In the example you have provided, if I am not mistaken, you solve difference equation while simultaneously learning the parameters of the difference equation. In the previous issue #56, we could create a composite node that fused the ODE solution such that we could obtain parameter estimates of the underlying ODE. This approach worked because the ODE was solved using DifferentialEquations.jl and RxInfer could evaluate logpdf of Poisson likelihood using the observations and the transformed ODE solution. Nevertheless in the current scenario, there are no observations to evaluate logpdfs nor a solution by a solver. If you use an outside solver then RxInfer can learn the parameters with the trick in #56 , otherwise it is not straightforward since the aim in this scenario is to solve the difference equation using RxInfer while learning the parameters. Once we progress in that roadmap, one can solve these scenarios as well. I am not really familiar with Gen.jl so my understanding of what the example does might be completely off. If that is the case perhaps a bit more context would be appreciated.

That being said,It would still be possible to obtain a solution in the current state of RxInfer but the solution will be convoluted and not straightforward to explain non-developers. Perhaps, I or @bvdmitri can have a look how we can solve this problem once we have time. I hope this explanation helps.

[sdwfrost]

Thanks @ismailsenoz - the basic problem is that one has to infer the latent states of the model in addition to the underlying parameters - this greatly increases the number of parameters to be inferred (in the example I give, you have to infer l 'recovery' events). Mixing of MCMC approaches to this tends not to be very good. Given 'observed' latent variables, you can get a logpdf, but if you don't have that, then you can get an estimate of the logpdf using, say, a particle filter (see here. I don't see any hurdles to applying the same logpdf trick applied in the ODE example, but it would be great if the model didn't have to be fused into a single node. You can just replace the binomials with Poissons (although there's a risk of getting negative counts) if that would be a blocker. Pardon my ignorance on the nuts and bolts of RxInfer - I still have to read @bvdmitri thesis, which looks like a good place to get a foundational perspective.

[ismailsenoz]

struct BinomialRx <: DiscreteUnivariateDistribution
    I 
    S 
    β 
end

function BayesBase.logpdf(bin::BinomialRx, y)
    I = bin.I
    S = bin.S
    β = bin.β

   return logpdf(convert(ExponentialFamilyDistribution, Binomial(S, ifrac_nonlinearity(I, β))), Int(y))
end

@node BinomialRx Stochastic [ out, I, S, β]

ifrac_nonlinearity(I, β; c = 10.0, N = 1000, δt = 1.0) = 1 - exp(-β*c*I/N*δt)
s_nonlinearity(s, i) = s - i < 0 ? 0 : s - i
i_nonlinearity(I, infection, recovery) = I + infection - recovery < 0 ? 0 : I + infection - recovery 

Distributions.logpdf(::Uninformative, x) = 1
Base.length(::Uninformative) = 1
function BayesBase.insupport(::Uninformative, y)
    return true
end
γ = 0.25
δt = 1
rfrac = 1 - exp(-γ*δt)
@model function sir_markov_model(rfrac)
    I[1] ~ Poisson(100) ## DiscreteUniform(1, 100)
    β    ~ Beta(1, 20)
    S[1] := s_nonlinearity(1000, I[1])
        for i in 1:39
            infection[i] ~ BinomialRx(I[i], S[i], β)
            recovery[i] ~ BinomialRx(I[i], S[i], rfrac)
            S[i+1] := s_nonlinearity(S[i], infection[i])
            I[i+1] := i_nonlinearity(I[i], infection[i], recovery[i])
        end
    S[40] ~ Uninformative()
    I[40] ~ Uninformative()
end;

I think a model of this sort could work and can support a mechanism where the simulation of the underlying SDE is also done purely with RxInfer and you will have access to recovery and infection. That being said, even with the current release you will not be able to run inference on this model without changing the internals. If you want to solve the SDE using DifferentialEquations.jl and use RxInfer to estimate the parameters, approach in #56 would certainly be possible. I will try to see if this model can work and keep you updated.

[sdwfrost]

Thanks @ismailsenoz ! If you could get the model working within RxInfer, that would be great, as many PPLs struggle with this.