Gaining speed in sampling an ODE

I was really pleased to see DifferentialEquations in pymc3. So, I started working on one problem of mine, where I could need this. Because currently, I do ML estimates, but going Bayesian would be great. I’m already very greatfull for the support I got from @dpananos and @michaelosthege (Return value for 2n-dimensional ODE system). Thanks!

My main problem now is, that pymc3 is much too slow when sampling. Here is an example notebook:

It is quite lengthy, as it is the actual task I am working on. I simulate some data and perform parameter recovery. The original data I work on has the same structure.

I infer parameters for two different models. The first model (A) is an n-dimensional system of ODEs and there are only 2 free parameters. This model seems still doable and I get reasonable results. Still, it is quite slow. And posterior analysis, things like traceplot are very slow as well.

However, the second model (D) which is a 2*n-dimensional ODE system and has 5 free parameters is way too slow to finish in reasonable time on my machine. But this is the model I work on and ML estimates seem to be OK (not shown in the notebook).

I already hacked the DifferentialEquation class to use solve_ivp which is much faster than odeint in my case (this also inspired Solve_ivp for Differential Equation).

Is there anything else to speed up computations? Any advise would be very welcome.

I know this seems a bit silly to be saying on PyMC3’s discourse, but have you maybe thought about using Stan for their ODE capability? Their implementation absolutely crushes mine, and so if speed is a concern, that would be my recomendation. DifferentialEquation is still new functionality and there is a lot of work to be done.

1 Like