timing code for comparing pypy/numba

test_growcut.py

from timeit import timeit

try:
    import numpypy as np
except:
    import numpy as np

import math

def window_floor(idx, radius):
    if radius > idx:
        return 0
    else:
        return idx - radius


def window_ceil(idx, ceil, radius):
    if idx + radius > ceil:
        return ceil
    else:
        return idx + radius

def python_kernel(image, state, state_next, window_radius):
    changes = 0
    sqrt_3 = math.sqrt(3.0)

    height = image.shape[0]
    width = image.shape[1]

    for j in xrange(width):
        for i in xrange(height):

            winning_colony = state[i, j, 0]
            defense_strength = state[i, j, 1]

            for jj in xrange(window_floor(j, window_radius),
                             window_ceil(j+1, width, window_radius)):
                for ii in xrange(window_floor(i, window_radius),
                                 window_ceil(i+1, height, window_radius)):
                    if (ii == i and jj == j):
                        continue

                    d = image[i, j, 0] - image[ii, jj, 0]
                    s = d * d
                    for k in range(1, 3):
                        d = image[i, j, k] - image[ii, jj, k]
                        s += d * d
                    gval = 1.0 - math.sqrt(s)/sqrt_3

                    attack_strength = gval * state[ii, jj, 1]

                    if attack_strength > defense_strength:
                        defense_strength = attack_strength
                        winning_colony = state[ii, jj, 0]
                        changes += 1

            state_next[i, j, 0] = winning_colony
            state_next[i, j, 1] = defense_strength

    return changes

try:
    from numba import autojit
    window_ceil  = autojit()(window_ceil)
    window_floor = autojit()(window_floor)
    kernel = autojit()(python_kernel)
    name = "numba kernel"
except:
    kernel = python_kernel
    name = "pypy kernel"

def time_kernel(k_name, N, dtype):
    """Timeit Helper Function (GrowCut Functions)

    Parameters
    ----------
    k_name : str
        Name of the GrowCut kernel in the __main__ namespace to be timed.
    N : int
        Size of image arrays to construct and pass to function.
    dtype : np.dtype
        Type of arrays to construct and pass to function.

    Returns
    -------
    func_time : float
        Average execution time over 3 trials
    """

    import __main__
    image = np.zeros((N, N, 3), dtype=dtype)
    state = np.zeros((N, N, 2), dtype=dtype)
    state_next = np.zeros((N, N, 2), dtype=dtype)


    # colony 1 is strength 1 at position 0,0
    # colony 0 is strength 0 at all other positions
    state[0, 0, 0] = 1
    state[0, 0, 1] = 1

    __main__.image = image
    __main__.state = state
    __main__.state_next = state_next

    trials = 3

    return timeit(stmt="kernel(__main__.image, __main__.state, __main__.state_next, 10)",
                  setup="from __main__ import %s as kernel; import __main__" % k_name,
                  number=trials)/trials

# throwaway
time_kernel('kernel', 50, np.float64)

print name
print '50x50', time_kernel('kernel', 50, np.float64)
print '100x100', time_kernel('kernel', 100, np.float64)
print '500x500', time_kernel('kernel', 500, np.float64)

On an Intel Core i7

~/sandbox ❯❯❯ pypy --version ⏎
Python 2.7.3 (b9c3566aa017, May 09 2013, 15:27:31)
[PyPy 2.0.0 with GCC 4.2.1 Compatible Apple LLVM 4.2 (clang-425.0.27)]

~/sandbox ❯❯❯ pypy test_growcut.py ⏎
pypy kernel
50x50 0.106865644455
100x100 0.227809031804
500x500 6.10808030764
~/sandbox ❯❯❯ ~/sandbox/anaconda/bin/ipython test_growcut.py
numba kernel
50x50 0.00457604726156
100x100 0.0203083356222
500x500 0.560228983561