Code Snippets for Medium blog post

numba_jit_example.py

from numba import jit


@jit
def jitted_func(x, y, z, overlay=False):
    "Predicts a quantity at times = 0, 1, ... 14"
    out = np.zeros((x.shape[0], 15))
    for t in range(15):
        out[:, t] = t * x ** 2 + y - 2 * z - 2 * t
    adj = 1.5 if overlay else 1.0
    return adj * out

  
# create some artificial inputs
n = 25000
u = np.random.random(n)
x = np.random.poisson(lam=5, size=n)
y, z = np.random.normal(size=(n, 2)).T


%%timeit -n 100
_ = predict_over_time(x, y, z) # 100 loops, best of 3: 3.28 ms per loop


%%timeit -n 100
_ = jitted_func(x, y, z) # 100 loops, best of 3: 2.27 ms per loop