This is one of the 100 recipes of the IPython Cookbook, the definitive guide to high-performance scientific computing and data science in Python.

5.1. Accelerating pure Python code with Numba and Just-In-Time compilation

In this example, we first write a pure Python version of a function that generates a Mandelbrot fractal. Then, we use Numba to compile it dynamically to native code.

In [ ]:
import numpy as np

We initialize the simulation and generate the grid in the complex plane.

In [ ]:
size = 200
iterations = 100

Pure Python version

The following function generates the fractal.

In [ ]:
def mandelbrot_python(m, size, iterations):
    for i in range(size):
        for j in range(size):
            c = -2 + 3./size*j + 1j*(1.5-3./size*i)
            z = 0
            for n in range(iterations):
                if np.abs(z) <= 10:
                    z = z*z + c
                    m[i, j] = n
                else:
                    break
In [ ]:
m = np.zeros((size, size))
mandelbrot_python(m, size, iterations)
In [ ]:
import matplotlib.pyplot as plt
%matplotlib inline
plt.imshow(np.log(m), cmap=plt.cm.hot,);
plt.xticks([]); plt.yticks([]);
In [ ]:
%%timeit m = np.zeros((size, size))
mandelbrot_python(m, size, iterations)

Numba version

We first import Numba.

In [ ]:
import numba
from numba import jit, complex128

Now, we just add the @jit decorator to the exact same function.

In [ ]:
@jit(locals=dict(c=complex128, z=complex128))
def mandelbrot_numba(m, size, iterations):
    for i in range(size):
        for j in range(size):
            c = -2 + 3./size*j + 1j*(1.5-3./size*i)
            z = 0
            for n in range(iterations):
                if abs(z) <= 10:
                    z = z*z + c
                    m[i, j] = n
                else:
                    break
In [ ]:
m = np.zeros((size, size))
mandelbrot_numba(m, size, iterations)
In [ ]:
%%timeit m = np.zeros((size, size))
mandelbrot_numba(m, size, iterations)

The Numba version is 250 times faster than the pure Python version here!

You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).

IPython Cookbook, by Cyrille Rossant, Packt Publishing, 2014 (500 pages).