How can I use Cython well to solve a differential equation faster?

Question

I would like to lower the time Scipy's odeint takes for solving a differential equation.

To practice, I used the example covered in Python in scientific computations as template. Because odeint takes a function f as argument, I wrote this function as a statically typed Cython version and hoped the running time of odeint would decrease significantly.

The function f is contained in file called ode.pyx as follows:

import numpy as np
cimport numpy as np
from libc.math cimport sin, cos

def f(y, t, params):
  cdef double theta = y[0], omega = y[1]
  cdef double Q = params[0], d = params[1], Omega = params[2]
  cdef double derivs[2]
  derivs[0] = omega
  derivs[1] = -omega/Q + np.sin(theta) + d*np.cos(Omega*t)
  return derivs

def fCMath(y, double t, params):
  cdef double theta = y[0], omega = y[1]
  cdef double Q = params[0], d = params[1], Omega = params[2]
  cdef double derivs[2]
  derivs[0] = omega
  derivs[1] = -omega/Q + sin(theta) + d*cos(Omega*t)
  return derivs

I then create a file setup.py to complie the function:

from distutils.core import setup
from Cython.Build import cythonize

setup(ext_modules=cythonize('ode.pyx'))

The script solving the differential equation (also containing the Python version of f) is called solveODE.py and looks as:

import ode
import numpy as np
from scipy.integrate import odeint
import time

def f(y, t, params):
    theta, omega = y
    Q, d, Omega = params
    derivs = [omega,
             -omega/Q + np.sin(theta) + d*np.cos(Omega*t)]
    return derivs

params = np.array([2.0, 1.5, 0.65])
y0 = np.array([0.0, 0.0])
t = np.arange(0., 200., 0.05)

start_time = time.time()
odeint(f, y0, t, args=(params,))
print("The Python Code took: %.6s seconds" % (time.time() - start_time))

start_time = time.time()
odeint(ode.f, y0, t, args=(params,))
print("The Cython Code took: %.6s seconds ---" % (time.time() - start_time))

start_time = time.time()
odeint(ode.fCMath, y0, t, args=(params,))
print("The Cython Code incorpoarting two of DavidW_s suggestions took: %.6s seconds ---" % (time.time() - start_time))

I then run:

python setup.py build_ext --inplace
python solveODE.py 

in the terminal.

The time for the python version is approximately 0.055 seconds, whilst the Cython version takes roughly 0.04 seconds.

Does somebody have a recommendation to improve on my attempt of solving the differential equation, preferably without tinkering with the odeint routine itself, with Cython?

Edit

I incorporated DavidW's suggestion in the two files ode.pyx and solveODE.py It took only roughly 0.015 seconds to run the code with these suggestions.

Answer 1

The easiest change to make (which will probably gain you a lot) is to use the C math library sin and cos for operations on single numbers instead of number. The call to numpy and the time spent working out that it isn't an array is fairly costly.

from libc.math cimport sin, cos

    # later
    -omega/Q + sin(theta) + d*cos(Omega*t)

I'd be tempted to assign a type to the input d (none of the other inputs are easily typed without changing the interface):

def f(y, double t, params):

I think I'd also just return a list like you do in your Python version. I don't think you gain a lot by using a C array.

Answer 2

tldr; use numba.jit for 3x speedup...

I don't have much experience with cython, but my machine seems to get similar computation times for your strictly python version, so we should be able to compare roughly apples to apples. I used numba to compile the function f (which I re-wrote slightly to make it play nicer with the compiler).

def f(y, t, params):
    return np.array([y[1], -y[1]/params[0] + np.sin(y[0]) + params[1]*np.cos(params[2]*t)])

numba_f = numba.jit(f)

dropping in numba_f in place of your ode.f gives me this output...

The Python Code took: 0.0468 seconds
The Numba Code took: 0.0155 seconds

I then wondered if I could duplicate odeint and also compile with numba to speed things up even further... (I could not)

Here is my Runge-Kutta numerical differential equation integrator:

#function f is provided inline (not as an arg)
def runge_kutta(y0, steps, dt, args=()): #improvement on euler's method. *note: time steps given in number of steps and dt
    Y = np.empty([steps,y0.shape[0]])
    Y[0] = y0
    t = 0
    n = 0
    for n in range(steps-1):
        #calculate coeficients
        k1 = f(Y[n], t, args) #(euler's method coeficient) beginning of interval
        k2 = f(Y[n] + (dt * k1 / 2), t + (dt/2), args) #interval midpoint A
        k3 = f(Y[n] + (dt * k2 / 2), t + (dt/2), args) #interval midpoint B
        k4 = f(Y[n] + dt * k3, t + dt, args) #interval end point

        Y[n + 1] = Y[n] + (dt/6) * (k1 + 2*k2 + 2*k3 + k4) #calculate Y(n+1)
        t += dt #calculate t(n+1)
    return Y

naive looping functions are typically the fastest once compiled, although this could probably be re-structured for a little better speed. I should note, this gives a different answer than odeint, deviating by as much as .001 after around 2000 steps, and is completely different after 3000. For the numba version of the function, I simply replaced f with numba_f, and added the compilation with @numba.jit as a decorator. In this case, as expected the pure python version is very slow, but the numba version is not any faster than the numba with odeint (again, ymmv).

using custom integrator
The Python Code took: 0.2340 seconds
The Numba Code took: 0.0156 seconds

Here's an example of compiling ahead of time. I don't have the necessary toolchain on this computer to compile, and I don't have admin to install it, so this gives me an error that I don't have the required compiler, but it should work otherwise.

import numpy as np
from numba.pycc import CC

cc = CC('diffeq')

@cc.export('func', 'f8[:](f8[:], f8, f8[:])')
def func(y, t, params):
    return np.array([y[1], -y[1]/params[0] + np.sin(y[0]) + params[1]*np.cos(params[2]*t)])

cc.compile()

Answer 3

If others answer this question using other modules, I might as well chime in:

I am the author of JiTCODE, which accepts an ODE written in SymPy symbols and then converts this ODE to C code for a Python module, compiles this C code, loads the result and uses this as a derivative for SciPy’s ODE. Your example translated to JiTCODE looks like this:

from jitcode import jitcode, provide_basic_symbols
import numpy as np
from sympy import sin, cos
import time

Q = 2.0
d = 1.5
Ω = 0.65

t, y = provide_basic_symbols()

f = [
    y(1),
    -y(1)/Q + sin(y(0)) + d*cos(Ω*t)
    ]

initial_state = np.array([0.0,0.0])

ODE = jitcode(f)
ODE.set_integrator("lsoda")
ODE.set_initial_value(initial_state,0.0)

start_time = time.time()
data = np.vstack(ODE.integrate(T) for T in np.arange(0.05, 200., 0.05))
end_time = time.time()
print("JiTCODE took: %.6s seconds" % (end_time - start_time))

This takes 0.11 seconds, which is horribly slow compared to the solutions based on odeint, but this is not due to the actual integration but the way the results are handled: While odeint directly creates an array efficiently internally, this is done via Python here. Depending on what you do, this may be a crucial disadvantage, but this quickly becomes irrelevant for a coarser sampling or larger differential equations.

So, let’s remove the data collection and just look at the integration, by replacing the last lines with the following:

ODE = jitcode(f)
ODE.set_integrator("lsoda", max_step=0.05, nsteps=1e10)
ODE.set_initial_value(initial_state,0.0)

start_time = time.time()
ODE.integrate(200.0)
end_time = time.time()
print("JiTCODE took: %.6s seconds" % (end_time - start_time))

Note that I set max_step=0.05 to force the integrator to make at least as many steps as in your example and ensure that the only difference is that the results of the integration are not stored to some array. This runs in 0.010 seconds.

Answer 4

NumbaLSODA takes 0.00088 seconds (17x faster than Cython).

from NumbaLSODA import lsoda_sig, lsoda
import numba as nb
import numpy as np
import time

@nb.cfunc(lsoda_sig)
def f(t, y_, dy, p_):
    p = nb.carray(p_, (3,))
    y = nb.carray(y_, (2,))
    theta, omega = y
    Q, d, Omega = p
    dy[0] = omega
    dy[1] = -omega/Q + np.sin(theta) + d*np.cos(Omega*t)

funcptr = f.address # address to ODE function
y0 = np.array([0.0, 0.0])
data = np.array([2.0, 1.5, 0.65])
t = np.arange(0., 200., 0.05)

start_time = time.time()
usol, success = lsoda(funcptr, y0, t, data = data)
print("NumbaLSODA took: %.8s seconds ---" % (time.time() - start_time))

result

NumbaLSODA took: 0.000880 seconds ---