Hello everyone,
First time posting here. I provide a lengthy intro to my research problem with hopefully enough details, aiming to discuss scalable solvers for Nonlinear systems of ODES. The question is at the end of this message.
My work is currently involved with a Bayesian continuous time deterministic SIR compartmental model with a time-varying transmission rate \beta_t following a random walk, on daily COVID-19 mortality data. The SIR ODE function keeps information about the three compartments (S,I,R) and a dummy compartment C is used to store the daily counts of infected individuals, for further use.
Letting N be the number of states and P be the number of model parameters, we have N>>P.
I initially coded the model in Stan, and performed MCMC with NUTS. I have used a Forward Euler ODE solver and the analysis is successfully executed (i.e. convergence diagnostics and posterior outputs are reasonable) in about 4 hrs under the setting: 395 time points, chains = 2, warmup = 500, main iters = 500), using rstan 2.21.2 on a Windows 10 machine (Intel Core i7, 16GB RAM).
I have successfully replicated the analysis using the trapezoidal rule for solving the ODE, and the same sampler takes about 8.3 hrs. The posterior outputs are nearly indistinguishable with the Euler implementation.
I wish to extend this work to an Multi-type (age-structured) SIR model with a time-varying infection rate. Eg, for 4 age groups, we end up with 16 states for this Multi-type SIR model. The greater the length of the time series, the bigger the dimension of the parameter space becomes, i.e. for 300 time points I have 300 transmission rate params, plus the infection rate parameter.
My testing attempt to execute NUTS with 500+500 iterations and the Euler solver takes 8 days of CPU time for a time series of 90 days (3 months).
I wish to use a scalable ODE solver and later also extend to more complex compartmental models like SEEIIR.
My search lead to the use of the adjoint solver in Stan. I’ve attempted analyzing in Stan under the aforementioned setting for the SIR, however the posterior outputs are not reasonable and MCMC diagnostics fail. To that end, I previously posted a related question on the Stan discourse forum. The adjoint solver appears to be at an experimental stage, and I’m currently looking for alternative ways to implement my analysis. I do not want to resort to discrete time stochastic compartmental models just yet.
My search in the last couple of days lead to the DifferentialEquations.jl suite and the web page of J. Storopoli describing Bayesian ODE models (https://storopoli.io/Bayesian-Julia/pages/12_epi_models)/).
I was wondering if the adjoint solver has been successfully tested for such compartmental models in Julia. Or if a Forward solver would be fast enough under high-dimensional parameter spaces and longer time series to allow the MCMC sampler to run within a couple of days for such Multi-type models.
In the end, I wish to create an R package performing MCMC of such epidemic models.
Thanks.
