Optimizing the Cython code for solving system of ODEs

Question

I am trying to optimize this code for solving a system of ODE. It seems Cython does not speed up the code compared to code using numpy.

This the python code using numpy:

"""Solve system of ODEs."""

def solver(ode_sys, I, t, integration_method):

    N = len(t)-1
    u = np.zeros((N+1, len(I)))
    u[0, :] = I
    dt = t[1] - t[0]

    for n in range(N):
        u[n+1, :] = integration_method(u[n, :], t[n], dt, n, ode_sys)
    return u, t


def RK2(u, t, dt, n, ode_sys):

    K1 = dt * ode_sys(u, t)
    K2 = dt * ode_sys(u + 0.5 * K1, t + 0.5 * dt)
    unew = u + K2
    return unew


def problem1(u, t):
    return -u + 1.0

from numpy import exp

def problem2(u, t):
    return - u + exp(-2.0 * t)

to run the code:

def run(ode_sys, N, nperiods=40):
    I = np.ones(N)
    time_points = np.linspace(0, nperiods * 2 * np.pi, nperiods * 30 + 1)
    u, t = solver(ode_sys, I, time_points, RK2)

timing

%timeit run(problem1, 1000, 1000)
%timeit run(problem2, 100, 500)

418 ms ± 22.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
233 ms ± 79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

My Cython version: I directly call the functions inside the module to get more speedup

%%cython --annotate
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.float_t DT

cdef extern from "math.h":
    double exp(double)

@cython.boundscheck(False) # turn off bounds checking for this func.
@cython.wraparound(False)  # Deactivate negative indexing.    
cpdef solver(ode_sys, np.ndarray[DT, ndim=1, negative_indices=False, mode='c'] I, 
             np.ndarray[DT, ndim=1, negative_indices=False, mode='c'] t, 
             integration_method):

    cdef int N = len(t)-1
    cdef np.ndarray[DT, ndim=2, negative_indices=False, 
                    mode='c'] u = np.zeros((N+1, len(I)))
    u[0, :] = I
    cdef double dt = t[1] - t[0]
    cdef int n
    for n in range(N):
        u[n+1, :] = RK2(u[n, :], t[n], dt, n, ode_sys)
    return u, t


def RK2(np.ndarray[DT, ndim=1, negative_indices=False, mode='c'] u, 
        double t, double dt, int n, ode_sys):

    cdef np.ndarray[DT, ndim=1, negative_indices=False, mode='c'] K1, K2, unew
    K1 = dt * problem1(u, t)
    K2 = dt * problem1(u + 0.5 * K1, t + 0.5 * dt)
    unew = u + K2
    return unew

cdef problem1(np.ndarray[DT, ndim=1, negative_indices=False, mode='c'] u, double t):
    return -u + 1.0

cdef problem2(np.ndarray[DT, ndim=1, negative_indices=False, mode='c'] u, double t):
    return - u + exp(-2.0 * t)

timing:

%timeit run(problem1, 1000, 1000) # note that to check promlem2 I need to change it inside the modules

424 ms ± 8.23 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Thanks in advance for any guide.

Answer 1

I have not checked the outputs of the run commands. I know two things to study your code performance.

Use the cython -a command (which I see in your code)

In jupyter notebook after executing this command, it will display below your code with lines in white and yellow. The darker the line is, the more overhead there will be because of python interactions. You can checkout this cython documentation page. In your code, lines involving numpy arrays are involved. These operations are calling numpy's python functions. I did an attempt to tackle this problem.

Profile it.

You can check this page. It shows how to profile your code with pstats and cProfile.

My solution (not the best/definite)

On my computer, the python version gives:

%timeit run(problem1, 1000, 1000)
564 ms ± 29.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

And the cython one:

%timeit run(problem1, 1000, 1000)
316 ms ± 17.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Also I forgot the decorators u added. Still, this is not a big improvement...

My final word is minimize the use/call to python function, since it will cause some overhead. And use elementwise operations on numpy arrays.

Here is my solution, by defining some operations replacing numpy's:

%%cython --annotate
# cython: profile=False
# cython: wraparound=False, boundscheck=False
import numpy as np
cimport numpy as cnp
cnp.import_array()
cimport cython
ctypedef cnp.double_t DT

cdef extern from "math.h":
    double exp(double)

cdef cnp.ndarray add(cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] X1, cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] X2):
    """sum two numpy arrays elementwise"""
    cdef cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] res = X1
    cdef int n = X1.shape[0]
    cdef int i

    for i in range(n):
        res[i] += X2[i]
    return res

cdef cnp.ndarray constant_multiply(double c, cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] X):
    """product of a constants and an array"""
    cdef cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] res = X
    cdef int n = X.shape[0]
    cdef int i
    
    for i in range(n):
        res[i] *= c
    return res
 
cpdef tuple solver(ode_sys, cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] I, 
             cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] t, 
             integration_method):

    cdef int N = t.shape[0]-1
    cdef cnp.ndarray[DT, ndim=2, negative_indices=False, 
                    mode='c'] u = np.zeros((N+1, I.shape[0]), dtype=np.float64)
    u[0, :] = I
    cdef double dt = t[1] - t[0]
    cdef int n
    for n in range(N):
        u[n+1, :] = RK2(u[n, :], t[n], dt, n, ode_sys)
    return u, t


cpdef cnp.ndarray RK2(cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] u, 
        double t, double dt, int n, ode_sys):

    cdef cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] K1, K2, unew
    K1 = constant_multiply(dt, problem1(u, t)) #dt * problem1(u, t)
    K2 = constant_multiply(dt, problem1(
        add(u, constant_multiply(0.5, K1)), 
        t + 0.5 * dt)
    ) #dt * problem1(u + 0.5 * K1, t + 0.5 * dt)
    unew = add(u, K2) #u + K2
    return unew

cdef cnp.ndarray problem1(cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] u, double t):
    """element wise problem 1"""
    cdef cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] res = u
    cdef int n = u.shape[0]
    cdef int i

    for i in range(n):
        res[i] = 1.0 - u[i]
    return res
    #return -u + 1.0

cdef cnp.ndarray problem2(cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] u, double t):
    """element wise problem 2"""
    cdef cnp.ndarray[DT, ndim=1, negative_indices=False, mode='c'] res = u
    cdef int n = u.shape[0]
    cdef int i

    for i in range(n):
        res[i] = - u[i] + exp(-2.0 * t)
    return res    
    #return - u + exp(-2.0 * t)

What to try (not an exhaustive list)

If you want to stick to Cython, there are still things to explore:

• use memory views instead of cnp.arrays in the typing
• use pythran backend, this is fairly easy to setup. You can keep your Cython code, add a comment line on top and compile it with some additional arguments. This documentation page says it all.
• parallellize it, but you cannot parallelize loops that contains Python operations

Other Options

There are other options if you dont want to use Cython:

• numba, this solution is pretty good since you have a numpy implementation and parallelizing is fairly easy.
• code some of your routines to fortran/C and use them in your python scripts. There is f2py that can help you by importing fortran routines. I never tried to code/import C routines.