Chapter 5, example 5

Here we show how to use Cython with NumPy to accelerate Python algorithms operating on arrays. The code consists in simulating a stochastic process.



In [1]:

    
%load_ext cythonmagic



In [2]:

    
n = 10000

This function simulates a brownian motion with random, independent steps (up or down). We use a Python loop here.



In [3]:

    
def step():
    return sign(rand(1) - .5)

def sim1(n):
    x = zeros(n)
    dx = 1./n
    for i in xrange(n - 1):
        x[i+1] = x[i] + dx * step()
    return x



In [4]:

    
plot(sim1(10000))









    Out[4]:





[<matplotlib.lines.Line2D at 0x5a99b90>]

Let's evaluate the performance of this function.



In [5]:

    
%timeit sim1(n)









    



1 loops, best of 3: 216 ms per loop

Now we convert it in Cython. We specify the type of each variable, and we rewrite the step function in C so that it only uses pure C functions.



In [6]:

    
%%cython
import numpy as np
# We need to import NumPy with cimport
cimport numpy as np
DTYPE = np.double
# We define the type of the array for Cython
ctypedef np.double_t DTYPE_t

# We rewrite the step function in C using the standard library
from libc.stdlib cimport rand, RAND_MAX
from libc.math cimport round

cdef double step():
    return 2 * round(float(rand()) / RAND_MAX) - 1

def sim2(int n):
    # here we specify the type of each variable with cdef, including the type of the array
    # with a special syntax so that Cython knows at compile-time the number of dimensions of this array.
    cdef int i
    cdef double dx = 1. / n
    cdef np.ndarray[DTYPE_t, ndim=1] x = np.zeros(n, dtype=DTYPE)
    for i in range(n - 1):
        x[i+1] = x[i] + dx * step()
    return x



In [7]:

    
%timeit sim2(n)









    



1000 loops, best of 3: 900 us per loop

We achieve a 240x speed improvement here!