Optimization of Scientific Code with Cython: Ising Model

Python is quick and easy to code, but can be slow when doing intensive numerical operations. Translating code to Cython can be helpful, but in most cases requires a bit of trial and error to achieve the optimal result. Cython's tutorials contain a lot of information, but for iterative workflows like optimization with Cython, it's often useful to see it done "live".

For that reason, I decided to record some screencasts showing this iterative optimization process, using an Ising Model, as an example application.

When to use Cython

Before I get to the videos, I wanted to say a few words about when and why you might choose Cython.

With scientific Python code, before turning to Cython I'd suggest going as far as you can with vectorization. Vectorization involves the judicious use of built-in routines in NumPy, SciPy, Pandas, and other libraries to reduce the number of explicit for-loops in your code. It can work quite well in many situations, and doesn't require any sort of special compilation step in running your code. See my PyCon 2015 talk, Losing Your Loops for an intro to this approach.

When a problem cannot be easily solved using standard vectorization approaches, Cython is a good choice. Cython provides a bridge between Python and C code, and is quite mature and reliable: it forms the basis of much of the PyData/Scientific Python ecosystem (for example, Cython is used heavily in NumPy, SciPy, Pandas, Scikit-Learn, and many other packages).

Other approaches to optimizing Python code are available, but I tend not to use them as often as Cython:

The PyPy project is an alternative implementation of Python that avoids the slow loops of the default CPython implementation. Although it is quite promising, it currently does not support many of the core scientific packages, so is not the best choice for scientific code (though that has been changing).

Numba is a Python-to-LLVM converter which can often give you 100x speedups in Python by adding a simple compilation decorator to Python functions. For an example of Numba being used on code like that in this notebook, see Matthew Rocklin's post. Though it is convenient, Numba doesn't support all Python constructs yet, and can be difficult to optimize when the most straightforward approach fails.

Videos

Finally, here are the demos of using Cython to optimize an Ising model.

In [1]:
from IPython.display import YouTubeVideo

Part 1

In part 1, I write a function to evolve an Ising model in Python, along with some tools to visualize the resulting evolution:

In [2]:
YouTubeVideo('rN7g4gzO2sk', width=600, height=350)
Out[2]:

Part 2

In part 2, I go through the process of optimizing this function with Cython:

In [3]:
YouTubeVideo('LOzcSuw3yOY', width=600, height=350)
Out[3]:

The Code

Below you can find the code that I wrote in those two videos:

1. Simple Python Ising Model

Displaying an Ising Field

In [4]:
import numpy as np

def random_spin_field(N, M):
    return np.random.choice([-1, 1], size=(N, M))

random_spin_field(10, 10)
Out[4]:
array([[ 1, -1, -1,  1, -1, -1,  1, -1,  1,  1],
       [ 1, -1,  1, -1, -1,  1, -1, -1, -1, -1],
       [-1,  1,  1,  1, -1, -1,  1,  1,  1, -1],
       [ 1, -1, -1,  1, -1,  1, -1,  1,  1, -1],
       [-1,  1, -1, -1, -1,  1, -1, -1, -1, -1],
       [-1,  1,  1, -1,  1, -1,  1,  1, -1, -1],
       [ 1, -1,  1, -1,  1, -1, -1,  1,  1, -1],
       [-1,  1, -1, -1,  1, -1, -1, -1, -1, -1],
       [ 1, -1, -1, -1,  1,  1, -1,  1, -1,  1],
       [ 1,  1,  1, -1,  1, -1,  1,  1, -1, -1]])
In [5]:
# pip install pillow
from PIL import Image

def display_spin_field(field):
    return Image.fromarray(np.uint8((field + 1) * 0.5 * 255))  # 0 ... 255

display_spin_field(random_spin_field(200, 200))
Out[5]:

Implementing the Ising Model

In [6]:
def ising_step(field, beta=0.4):
    N, M = field.shape
    for n_offset in range(2):
        for m_offset in range(2):
            for n in range(n_offset, N, 2):
                for m in range(m_offset, M, 2):
                    _ising_update(field, n, m, beta)
    return field

def _ising_update(field, n, m, beta):
    total = 0
    N, M = field.shape
    for i in range(n-1, n+2):
        for j in range(m-1, m+2):
            if i == n and j == m:
                continue
            total += field[i % N, j % M]
    dE = 2 * field[n, m] * total
    if dE <= 0:
        field[n, m] *= -1
    elif np.exp(-dE * beta) > np.random.rand():
        field[n, m] *= -1
In [7]:
display_spin_field(ising_step(random_spin_field(200, 200)))
Out[7]:

Animating an Ising Sequence

Note that this requires a live kernel, and so will not show up on the blog or other static rendering:

In [8]:
from ipywidgets import interact

def display_ising_sequence(images):
    def _show(frame=(0, len(images) - 1)):
        return display_spin_field(images[frame])
    return interact(_show)
In [9]:
images = [random_spin_field(200, 200)]
for i in range(50):
    images.append(ising_step(images[-1].copy()))
display_ising_sequence(images);

The result will look something like this:

[Ising-Slider-1.gif]

2. Use Cython to speed it up

In [10]:
%load_ext Cython
In [11]:
%%cython

cimport cython

import numpy as np
cimport numpy as np

from libc.math cimport exp
from libc.stdlib cimport rand
cdef extern from "limits.h":
    int RAND_MAX


@cython.boundscheck(False)
@cython.wraparound(False)
def cy_ising_step(np.int64_t[:, :] field, float beta=0.4):
    cdef int N = field.shape[0]
    cdef int M = field.shape[1]
    cdef int n_offset, m_offset, n, m
    for n_offset in range(2):
        for m_offset in range(2):
            for n in range(n_offset, N, 2):
                for m in range(m_offset, M, 2):
                    _cy_ising_update(field, n, m, beta)
    return np.array(field)


@cython.boundscheck(False)
@cython.wraparound(False)
cdef _cy_ising_update(np.int64_t[:, :] field, int n, int m, float beta):
    cdef int total = 0
    cdef int N = field.shape[0]
    cdef int M = field.shape[1]
    cdef int i, j
    for i in range(n-1, n+2):
        for j in range(m-1, m+2):
            if i == n and j == m:
                continue
            total += field[i % N, j % M]
    cdef float dE = 2 * field[n, m] * total
    if dE <= 0:
        field[n, m] *= -1
    elif exp(-dE * beta) * RAND_MAX > rand():
        field[n, m] *= -1

Timing the result

In [12]:
field = random_spin_field(200, 200)
%timeit ising_step(field)
%timeit cy_ising_step(field)
317 ms ± 10.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
2.64 ms ± 132 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Visualizing the result

(Note that the following code requires a live kernel, and will not display within a static rendering such as the blog post):

In [13]:
images = [random_spin_field(500, 500)]
for i in range(100):
    images.append(cy_ising_step(images[-1].copy(), beta=0.4))
display_ising_sequence(images);

The result will look something like this:

[Ising-Slider-2.gif]

This post was written entirely in the Jupyter notebook. You can download this notebook, or see a static view on nbviewer.

Comments