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.
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:
YouTubeVideo('rN7g4gzO2sk', width=600, height=350)
Part 2¶
In part 2, I go through the process of optimizing this function with Cython:
YouTubeVideo('LOzcSuw3yOY', width=600, height=350)
The Code¶
Below you can find the code that I wrote in those two videos:
1. Simple Python Ising Model¶
Displaying an Ising Field¶
import numpy as np
def random_spin_field(N, M):
return np.random.choice([-1, 1], size=(N, M))
random_spin_field(10, 10)
# 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))
Implementing the Ising Model¶
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
display_spin_field(ising_step(random_spin_field(200, 200)))
Animating an Ising Sequence¶
Note that this requires a live kernel, and so will not show up on the blog or other static rendering:
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)
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:
2. Use Cython to speed it up¶
%load_ext Cython
%%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¶
field = random_spin_field(200, 200)
%timeit ising_step(field)
%timeit cy_ising_step(field)
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):
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:
This post was written entirely in the Jupyter notebook. You can download this notebook, or see a static view on nbviewer.