Here we show how to use Cython with NumPy to accelerate Python algorithms operating on arrays. The code consists in simulating a stochastic process.
%load_ext cythonmagic
n = 10000
This function simulates a brownian motion with random, independent steps (up or down). We use a Python loop here.
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
plot(sim1(10000))
[<matplotlib.lines.Line2D at 0x5a99b90>]
Let's evaluate the performance of this function.
%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.
%%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
%timeit sim2(n)
1000 loops, best of 3: 900 us per loop
We achieve a 240x speed improvement here!