Hello everyone,

Our work is currently involved with a Bayesian SIR model with a time-varying infection rate on daily COVID-19 mortality data. The SIR ODE function keeps information about the three compartments (S,I,R) and a dummy compartment is used to store the time-varying infection rate for further use.

Letting N be the number of states and P be the number of model parameters, we have N >> P. We have used the RK45 solver and the analysis is successfully executed (i.e. convergence diagnostics and posterior outputs are reasonable) in about 46 hours under the setting (395 time points, chains = 2, warmup = 2000, main iters = 5000), using rstan 2.21.2 on a Windows 10 machine (Intel Core i7, 16GB RAM).

We wish to extend this work to an age-structured SIR model with a time-varying infection rate. Acknowledging the computational limitations of the analysis with a RK45 solver, a preliminary run using the RK45 solver was performed with (chains = 2, warmup = 500 iters, main = 700 iters). After 2 weeks, the progress was at about 60%.

Previous discussions (link1, link2, link3) have pointed towards the Adjoint ODE solver, expected to scale better under the N >> P situation. Taking a step back, we moved to cmdstan 2.27.0 and attempted to implement the Bayesian SIR model (no stratification) with the Adjoint ODE solver [ode_adjoint_tol_ctl(), cmdstan version 2.27.0, cmdstanr 0.4.0.90, R version 4.1.0). If this works, then weāll move to the age-structured model.

While the adjoint solver results into a fast implementation (~45 minutes!), a very similar informational message to what @Teddy_Groves1 has previously reported here comes up and all iterations had divergences. Weāve tried different different initialization points, but the issue persists.

Ideas from the people who have worked on this feature (and anyone else familiar with this feature) on how to configure the adjoint solver would be much appreciated.

I attach part of the .R file related to MCMC sampling and part of the corresponding .stan program; additionally, the error messages I receive after a short test run, for illustration purposes, if this can help. Also tagging @wds15, @charlesm93, @Funko_Unko in this.

Thanks!

Messages during compilation:

Messages during MCMC sampling:

Test_Adjoint_v2_forum.stan (4.8 KB)

AdjointODE_test.R (3.2 KB)