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.