nano-benchmark.py

import matplotlib.pyplot as plt
import timeit
import hashlib
import gc

import numpy as np
import numexpr as ne
from a import a as pythran
from numba import vectorize, jit
from pythran import compile_pythrancode
from imp import load_dynamic

orders_n = [10**n for n in range(1, 5)]

def make_numpy(func):
    return 'lambda x, y: ' + func

def make_numexpr(func):
    return 'lambda x, y: ne.evaluate("' + func + '")'

def make_numba(func):
    return 'jit(' + make_numpy(func) + ')'

def make_pythran(func):
    m = hashlib.md5() ; m.update(func) ; name = m.hexdigest()
    code = '#pythran export foo(float64[][], float64[][])\ndef foo(x, y): return ' + func
    so = compile_pythrancode('pythran_' + name, code)
    globals()['pythran_kernel'] = load_dynamic('pythran_' + name, so).foo
    return 'pythran_kernel'


expressions = (
    'x + y',
    '.8 * x + .2 * y',
    'x[:] + y [:]',
    'x + 1',
    'x + y[0]',
    'x[1:-1] + y[1:-1]',
    'x[1:-1,1:-1] + y[1:-1,1:-1]',
    'x[1:,1:] + x[:-1,1:] + x[1:,:-1] + x[:-1,:-1]',
    '(x - y)**2 / (x + y)',
    '.1 * x**2 + .2 * y ** 2 + .3 * x * y + .4 * x + .5 * y + .6',
)

engines = (
    'numpy',
    'numexpr',
    'pythran',
    'numba',
)

for nth, expression in enumerate(expressions):
    for engine in engines:
        locals()[engine] = eval('eval(make_' + engine + '(expression))')

    timings = {f:[] for f in engines}

    for n in orders_n:
        for f in engines:
            A = np.random.rand(n,n)
            B = np.random.rand(n,n)
            try:
                timing = min(timeit.Timer('%s(A, B)' %f,
                                          'from __main__ import A, B, %s' %f)
                             .repeat(repeat=3, number=10))
            except TypeError:
                timing = 0

            timings[f].append(timing)
            del A
            del B
            gc.collect()

    def plot_figures(engines):
        fig = plt.figure(figsize=(8,8))
        for f in engines:
            plt.plot([i**2 for i in orders_n], [i/n for i,n in zip(timings[f], timings['numpy'])],
                                alpha=0.5, label='%s' %f, marker='o', lw=2)
        plt.legend(loc='upper left')
        #plt.xscale('log')
        #plt.yscale('log')
        plt.axis(ymax=2.0,ymin=0.)
        plt.grid()
        plt.xlabel('number of elements per matrix')
        plt.ylabel('normalized execution time (w/ numpy) lower is better')
        plt.title('Approaches to evaluate ' + expression)

        plt.savefig('plot{}.png'.format(nth))

    plot_figures(engines)