SciPy+Numba odeint vs Julia ODE vs NumbaLSODA: 50x performance difference on stiff ODE

a_stiff_ode_performance_python_julia.md

SciPy+Numba odeint vs Julia DifferentialEquations.jl vs NumbaLSODA Summary

All are solved at reltol=1e-3, abstol=1e-6 using the fastest ODE solver of the respective package for the given problem.

Absolute Performance Numbers:

• SciPy LSODA through odeint takes ~489μs
• SciPy LSODA through odeint with Numba takes ~257μs
• NumbaLSODA takes ~25μs
• DifferentialEquations.jl Rosenbrock23 takes ~9.2μs

Relative Performance Numbers:

• SciPy LSODA through odeint takes 53x as long
• SciPy LSODA through odeint with Numba takes 28x as long
• numbalsoda takes 2.7x as long

numbalsoda.py

from numbalsoda import lsoda_sig, lsoda
from numba import njit, cfunc
import numpy as np
import timeit

@cfunc(lsoda_sig)
def rhs(t, u, du, p):
    k1 = p[0]
    k2 = p[1]
    k3 = p[2]
    du[0] = -k1*u[0]+k3*u[1]*u[2]
    du[1] = k1*u[0]-k2*u[1]**2-k3*u[1]*u[2]
    du[2] = k2*u[1]**2
    
u0 = np.r_[1.0, 0.0, 0.0]
t = np.r_[0.0, 1e5]
p = np.r_[0.04, 3e7, 1e4]
rhs_address = rhs.address
usol, success = lsoda(rhs_address, u0, t, data=p, rtol=1.0e-3, atol=1.0e-6)
assert success

@njit
def time_func():
    usol, success = lsoda(rhs_address, u0, t, data=p, rtol=1.0e-3, atol=1.0e-6)
    assert success

time_func()
print('NumbaLSODA', timeit.Timer(time_func).timeit(number=1000)/1000, 'seconds')
print('NumbaLSODA', timeit.Timer(time_func).timeit(number=1000)/1000, 'seconds')

rober.jl

using OrdinaryDiffEq, StaticArrays, BenchmarkTools
function rober(u,p,t)
    y₁,y₂,y₃ = u
    k₁,k₂,k₃ = p
    du1 = -k₁*y₁+k₃*y₂*y₃
    du2 =  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃
    du3 =  k₂*y₂^2
    SA[du1,du2,du3]
end
prob = ODEProblem{false}(rober,SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4])
# Defaults to reltol=1e-3, abstol=1e-6
@btime sol = solve(prob,Rosenbrock23(),save_everystep=false) # 9.300 μs (26 allocations: 3.03 KiB)

rober_numba.py

from scipy.integrate import odeint
import numba
import timeit
def rober(u,t):
  k1 = 0.04
  k2 = 3e7
  k3 = 1e4
  y1, y2, y3 = u
  dy1 = -k1*y1+k3*y2*y3
  dy2 =  k1*y1-k2*y2*y2-k3*y2*y3
  dy3 =  k2*y2*y2
  return [dy1,dy2,dy3]

u0  = [1.0,0.0,0.0]
t = [0.0, 1e5]
numba_f = numba.jit(rober,nopython=True)
sol = odeint(numba_f,u0,t,rtol=1e-3,atol=1e-6)

def time_func():
    sol = odeint(numba_f,u0,t,rtol=1e-3,atol=1e-6)

timeit.Timer(time_func).timeit(number=100)/100 # 0.00025674299999081993 seconds
timeit.Timer(time_func).timeit(number=100)/100 # 0.0002520930000173394 seconds

rober_odeint.py

from scipy.integrate import odeint
import timeit

def rober(u,t):
  k1 = 0.04
  k2 = 3e7
  k3 = 1e4
  y1, y2, y3 = u
  dy1 = -k1*y1+k3*y2*y3
  dy2 =  k1*y1-k2*y2*y2-k3*y2*y3
  dy3 =  k2*y2*y2
  return [dy1,dy2,dy3]

u0  = [1.0,0.0,0.0]
t = [0.0, 1e5]
sol = odeint(rober,u0,t,rtol=1e-3,atol=1e-6)

def time_func():
    sol = odeint(rober,u0,t,rtol=1e-3,atol=1e-6)

timeit.Timer(time_func).timeit(number=100)/100 # 0.00048923599999398 seconds
timeit.Timer(time_func).timeit(number=100)/100 # 0.00047277800000301796  seconds

Thanks!