# Scalable ODE solver for Bayesian multi-type compartmental models
**URL:** https://discourse.julialang.org/t/scalable-ode-solver-for-bayesian-multi-type-compartmental-models/72845
**Category:** Biology, Health, and Medicine
**Tags:** question, ode, bayesian-inference, covid
**Created:** [December 9, 2021, 6:00pm UTC](https://discourse.julialang.org/t/scalable-ode-solver-for-bayesian-multi-type-compartmental-models/72845 "2021-12-09T18:00:07Z")
**Posts on this page:** 3
**Page:** 1
### Author:  [@bernadette-eu](https://discourse.julialang.org/u/bernadette-eu)
#### Post date: [December 9, 2021, 6:00pm UTC](https://discourse.julialang.org/t/scalable-ode-solver-for-bayesian-multi-type-compartmental-models/72845/1 "2021-12-09T18:00:07Z")
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](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.
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [December 10, 2021, 11:27am UTC](https://discourse.julialang.org/t/scalable-ode-solver-for-bayesian-multi-type-compartmental-models/72845/2 "2021-12-10T11:27:58Z")
> [@bernadette-eu](#):
>
> I was wondering if the adjoint solver has been successfully tested for such compartmental models in Julia.
Differentiation of compartmental models is done all of the time in [https://www.biorxiv.org/content/10.1101/2020.11.28.402297v1](https://www.biorxiv.org/content/10.1101/2020.11.28.402297v1). Adjoints are used in the automated discovery of epidemic models showcased in:
[ ](https://www.youtube.com/watch?v=FihLyzdjN_8)
That said,
> [@bernadette-eu](#):
>
> 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.
Forward-mode AD is actually more efficient than adjoint methods here. See:
> **[A Comparison of Automatic Differentiation and Continuous Sensitivity Analysis...](https://ieeexplore.ieee.org/abstract/document/9622796)**
>
> Derivatives of differential equation solutions are commonly for parameter estimation, fitting neural differential equations, and as model diagnostics. However, with a litany of choices and a Cartesian product of potential methods, it can be difficult...
> **[Differential Methods for Assessing Sensitivity in Biological Models](https://www.biorxiv.org/content/10.1101/2021.11.15.468697v1)**
>
> Differential sensitivity analysis is indispensable in fitting parameters, understanding uncertainty, and forecasting the results of both thought and lab experiments. Although there are many methods currently available for performing differential...
For benchmarks against Stan with Turing+DifferentialEquations.jl outperforming, see:
[https://benchmarks.sciml.ai/html/ParameterEstimation/DiffEqBayesLotkaVolterra.html](https://benchmarks.sciml.ai/html/ParameterEstimation/DiffEqBayesLotkaVolterra.html)
(shows Julia 5x faster than Stan)
[https://benchmarks.sciml.ai/html/ParameterEstimation/DiffEqBayesFitzHughNagumo.html](https://benchmarks.sciml.ai/html/ParameterEstimation/DiffEqBayesFitzHughNagumo.html)
(shows Julia 3x faster than Stan)
As for the instabilities, there’s a long discussion about what’s going wrong in Stan and the source of its instabilities on differential equation models:
>
and so that showcases some ways in which the Julia tools are more stable.
---
### Author:  [@bernadette-eu](https://discourse.julialang.org/u/bernadette-eu)
#### Post date: [December 10, 2021, 11:43am UTC](https://discourse.julialang.org/t/scalable-ode-solver-for-bayesian-multi-type-compartmental-models/72845/3 "2021-12-10T11:43:29Z")
Thank you for this material, Chris. I’ll go through it.
So, it looks like it’s worth investing time to learn Julia for the purpose of my work.