Show how to speed up scipy.integrate.odeint simply by decorating the right-hand side with numba's jit function

NumbaODEExample.ipynb

I'm relatively new to Python and Numba, but when I tried to run this example and compare the results of the integration, I get different values from the modified code. I wonder if odeint doesn't rely on the y0 passed to it to not change due to the evaluation of the RHS. The jit speedup is remarkable, but if I haven't made a mistake I don't think the resulting answer is correct.

Does this also work with integrate.ode (not integrate.odeint) ? I use ode.set_f_params to set some parameters for calculating the RHS, but then it fails with some TypeError complaining that the required argument is not found. Any ideas?

Update: also fails without setting extra parameters. I assume set_initial_value is enough to screw up numba functionality

IMHO, a much simpler solution that "always works" is to use the C-based integrators in PyDSTool or on github. It includes automatic code generation to turn your string-based declarations of ODEs into actual C code (no python callbacks) and executes much faster. You just need to have gcc/gfortran and SWIG installed, which is easy with the conda package manager (among other solutions).

@KaBrrrp, did you find a solution for integrate.ode?

Hello @moble, I've tried to implement numba with my ode function, and I'm having trouble because I am passing a dictionary and an custom object as arguments in the odeint function. I've never used numba before, but I tried your code and some simple for loops to understand the gist of it, but I've hit a wall.

Link to my question on stackoverflow

IMHO, a much simpler solution that "always works" is to use the C-based integrators in PyDSTool or on github. It includes automatic code generation to turn your string-based declarations of ODEs into actual C code (no python callbacks) and executes much faster. You just need to have gcc/gfortran and SWIG installed, which is easy with the conda package manager (among other solutions).

Yes, but to implement the same model in PyDSTool you need to write much more code. "Simple is better than complex"

this gist is old (2014) but still relevant, imo.

i currently am not able to not use ode with that same error "TypeError: not enough arguments: expected 2, got 1"

note there is a suggestion to bypass the arg error using a python wrapper function to the jited function < https://stackoverflow.com/questions/32744658/using-numba-jit-with-scipy-integrate-ode >

i was able to use jit with odeint, as mentioned here, but i had problem with the following statement:

" Fortunately, the first argument to odeint is an array that gets thrown away anyway, so we can just replace the values in that array and return it."

in that case I had different results (absurd results) when messing with the original y memory,
instead i made a scratch variable preallocated outside to work the dy

scratch_dN = np.zeros((5,)) #scratch memory allocated externally to dNi, reused between calls
...
sol = odeint(dNi, Ni0, t, args=(scratch_dN,) )
...

where dNi had been defined as:

def dNi(N,t,scratch_dN): #!!    
    ... 
    vector_dNi = scratch_dN #rename reuse memory
   ... 
    vector_dNi[0] = ...
    ...#up to vector_dNi0[5]
    ...
    return vector_dNi  

@iurisegtovich Yeah, I think something has changed internally with scipy. Also note that odeint and ode are actually considered the "Old API" now; the new one uses solve_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.

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)