jitspeed.py

import numpy as np

from numba import jit, typeof
from numba.compiler import compile_isolated, Flags

def run_func(func, args):
    import time
    clock = time.process_time
    func(*args)
    times = []
    for i in range(5):
        t1 = clock()
        dt = clock() - t1
        times.append(dt)
    print("%60s: %.3f ms" % (str(func), min(times) * 1000))


def py_loop(dep):
    nx,ny = dep.shape
    max_triangles = (ny-2)*(nx-2)*2*3 + 2*(ny-2+nx-2)
    x = np.zeros((max_triangles, 3), dtype=dep.dtype)
    y = x.copy()
    z = x.copy()
    k = 0

    for i in range(1, nx-1):
        rim, rip = i - .5, i + .5
        for j in range(1, ny-1):
            rjm, rjp = j - .5, j + .5
            dep00 = dep[i,j]
            for di, dj in ((1,0),(-1,0),(0,1),(0,-1)):
                dep01 = dep[i+di, j+dj]
                if dep01 > dep00:
                    if dj==0:
                        x01 = i + .5*di
                        x1 = (x01, x01, x01, x01)
                    else:
                        y01 = j + .5*dj
                        y1 = (y01, y01, y01, y01)
                        x1 = (rim, rip, rip, rim)
                    z1 = (dep00, dep00, dep01, dep01)

                    x[k,:], y[k,:], z[k,:] = x1[:-1], y1[:-1], z1[:-1]
                    k += 1
                    x[k,:], y[k,:], z[k,:] = x1[1:], y1[1:], z1[1:]
                    k += 1

            x1 = (rim, rip, rip, rim)
            y1 = (rjp, rjp, rjm, rjm)
            z1 = (dep00,dep00,dep00,dep00)

            x[k,:], y[k,:], z[k,:] = x1[:-1], y1[:-1], z1[:-1]
            k += 1
            x[k,:], y[k,:], z[k,:] = x1[1:], y1[1:], z1[1:]
            k += 1

    return k, x, y, z


def compile_func(func, sig):
    if 1:
        flags = Flags()
        flags.set("nrt")
        flags.set("no_compile")
        compile_isolated(func, sig, flags=flags)
    else:
        jit(sig, nopython=True)(func)


dep = np.random.rand(2, 4)
sig = (typeof(dep),)

run_func(compile_func, (py_loop, sig,))