Last active
July 6, 2021 18:32
-
-
Save moble/3aa44230256b66956587 to your computer and use it in GitHub Desktop.
Show how to speed up scipy.integrate.odeint simply by decorating the right-hand side with numba's jit function
I wrote a wrapper to LSODA which has no overhead: https://github.com/Nicholaswogan/NumbaLSODA . During an ODE solve, the python interpreter is never used, so everything is fast for small problems:
from NumbaLSODA import lsoda_sig, lsoda
import numba as nb
@nb.cfunc(lsoda_sig)
def RHS_nb(t, y, dy, p):
dy[0],dy[1] = t*y[1],y[0]
funcptr = RHS_nb.address
@nb.njit()
def test():
sol, success = lsoda(funcptr, y0_, t)
y0_ = np.array(y0)
%timeit test()
result is
26.2 µs ± 342 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
@iurisegtovich Yeah, I think something has changed internally with scipy. Also note that
odeint
andode
are actually considered the "Old API" now; the new one usessolve_ivp
and friends — but these examples about 1,000 times slower with the new API!!! (I can imagine it's all in the overhead.)Anyway, your method is a good way to go. Numba also supports
jitclass
now, so you could also pass a more complicated object as one of the args, with all sorts of fancy capabilities.But maybe more importantly, python itself has sped up significantly, so that even using the naive approach in this notebook gives nearly the same speed as when using numba. Obviously, really complicated functions will still benefit from numba, but in this example numba is actually a little bit slower in my tests.