Gaining speed in sampling an ODE

You’re right: the context switch in the iterations of the ODE solver is a huge problem.

There are a few things that are not too difficult and could lead to some acceleration:

  • numba-compiled functions of sympy-derived augment_system (instead of theano)
  • adjusting absolute/relative tolerances for odeint
  • option to turn off sensitivities & using a gradient-free sampler (only small & easy models)

And then there’s a the optimal solution that is much faster, but for which we need more help to implement:

  1. make cross-OS conda-installable package that wraps sundials
  2. sympy analysis to get augment_system
  3. numba-compile augment_system generated code in a way that it acts directly on sundials data types

The optimal strategy actually avoids all the context switches and also avoids copying data around all the time. @aseyboldt has a proof of concept already, but point 1. as well as a clean object-oriented implementation are still under construction.
Based on Adrians work, I managed to build an entry point for the sympy-based analysis of a user-provided ODE system.
See here:

2 Likes