numba speeds up python code 30x to 50x in Ubuntu 14.04!

test-numba.py

# Albert Cardona 2015-06-02
# Test the speed up provided by just-in-time compilation of python code
# using the numba JIT framework.
# The operation to test is a trivial one: the difference of the logarithm,
# pairwise, in two arrays of randomly generated floating-point numbers.
# For a far more realistic test involving a python implementation of
# the non-uniform fast fourier transform, see:
# https://jakevdp.github.io/blog/2015/02/24/optimizing-python-with-numpy-and-numba/
#
# To resolve issues with numba compilation (mostly type mismatches), see:
# http://numba.pydata.org/numba-doc/dev/user/troubleshoot.html  
# 
#
# Instructions for installing numba 0.18.2 in a python 2.7.6 environment
# in Ubuntu 14.04 LTS, requiring LLVM-3.5:
# (Assumes you have python 2.7 and a python environment setup named "env".)
#
#   # Remove LLVM 3.4
#   $ sudo apt-get remove llvm
#   $ sudo apt-get install llvm-3.5
#   # Make LLVM-3.5 the default
#   $ sudo ln -s /usr/bin/llvm-config-3.5 /usr/local/bin/llvm-config
#
#   # Needed to overcome issue with llvmlite (ubuntu's package of llvm-3.5
#   # is older than llvmlite's ability to work with 3.5, and there are some differences).
#   # See: https://github.com/numba/llvmlite/issues/27
#   $ sudo apt-get install libedit-dev
#
#   $ workon env
#
#   # Necessary for python versions prior to 3.4
#   $ pip install enum34
#
#   # LLVM python wrapper created by the numba dev team
#   $ pip install llvmlite
#
#   # Necessary for numba to run (but not for compiling)
#   $ pip install funcsigs
#
#   $ pip install numba
#

from __future__ import division, print_function
import warnings
import numpy as np
import numba
import time

# With numba.
# nopython=True means an error will be raised when compilation is not possible
# Notice we can't create arrays within: have to pass them as argument
@numba.jit(nopython=True)
def diffLogsNumba(img1, img2, diff):
    # img1.shape(0) is the same as len(img1), both work
    # range or xrange, both work at the same speed
    for i in xrange(img1.shape[0]):
        # Could use math.log as well, both work at the same speed
        diff[i] = np.log(img1[i]) - np.log(img2[i])

# Plain python
def diffLogs(img1, img2, diff):
    for i in xrange(img1.shape[0]):
        diff[i] = np.log(img1[i]) - np.log(img2[i])

def timeIt(fn, img1, img2, diff):
    t0 = time.time()
    r = fn(img1, img2, diff)
    print("Elapsed time: %s seconds" % (time.time() - t0))
    return r

def test(array_size):
    print("###\nTesting for array size of %s" % array_size)
    rng = np.random.RandomState(0)
    # Create two arrays of identical dimensions, filled with random numbers
    img1 = rng.rand(array_size)
    img2 = rng.rand(len(img1))
    # Create the output array
    diff = np.zeros(len(img1), dtype=np.float64)

    print("Numba (cold):")
    timeIt(diffLogsNumba, img1, img2, diff)
    print("Numba (JIT warmed up):")
    timeIt(diffLogsNumba, img1, img2, diff)
    print("Plain python:")
    timeIt(diffLogs, img1, img2, diff)


# Results

"""

test(512 * 512)

###
Testing for array size of 262144
Numba (cold):
Elapsed time: 0.0487880706787 seconds
Numba (JIT warmed up):
Elapsed time: 0.0172371864319 seconds
Plain python:
Elapsed time: 0.532037973404 seconds


test(4096 * 4096)

###
Testing for array size of 16777216
Numba (cold):
Elapsed time: 1.12040400505 seconds
Numba (JIT warmed up):
Elapsed time: 1.09282398224 seconds
Plain python:
Elapsed time: 33.5452198982 seconds

"""