<jupyter_start><jupyter_text>Just-in-time compilation (JIT)====For programmer productivity, it often makes sense to code the majority of your application in a high-level language such as Python and only optimize code bottlenecks identified by profiling. One way to speed up these bottlenecks is to compile the code to machine executables, often via an intermediate C or C-like stage. There are two common approaches to compiling Python code - using a Just-In-Time (JIT) compiler and using Cython for Ahead of Time (AOT) compilation.This notebook mostly illustrates the JIT approach. **References**- [Numba](http://numba.pydata.org)- [The need for speed without bothering too much: An introduction to numba](http://nbviewer.jupyter.org/github/akittas/presentations/blob/master/pythess/numba/numba.ipynb?utm_source=newsletter_mailer&utm_medium=email&utm_campaign=weekly)<jupyter_code>%matplotlib inline
import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>**Utility function for timing functions**We write decorators to help in timing as an alternative to `timeit`.<jupyter_code>import time
from numpy.testing import assert_almost_equal
def timer(f, *args, **kwargs):
    start = time.time()
    ans = f(*args, **kwargs)
    return ans, time.time() - start
def report(fs, *args, **kwargs):
    ans, t = timer(fs[0], *args, **kwargs)
    print('%s: %.1f' % (fs[0].__name__, 1.0))  
    for f in fs[1:]:
        ans_, t_ = timer(f, *args, **kwargs)
        print('%s: %.1f' % (f.__name__, t/t_))<jupyter_output><empty_output><jupyter_text>Using `numexpr`----One of the simplest approaches is to use [`numexpr`](https://github.com/pydata/numexpr) which takes a `numpy` expression and compiles a more efficient version of the `numpy` expression written as a string. If there is a simple expression that is taking too long, this is a good choice due to its simplicity. However, it is quite limited.<jupyter_code>import numpy as np
a = np.random.random(int(1e6))
b = np.random.random(int(1e6))
c = np.random.random(int(1e6))
%timeit -r3 -n3 b**2 - 4*a*c
import numexpr as ne # For simple numerical calculations
%timeit -r3 -n3 ne.evaluate('b**2 - 4*a*c')<jupyter_output>1.41 ms ± 632 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Using `numba`----When it works, the JIT `numba` can speed up Python code tremendously with minimal effort. [Documentation for `numba`](http://numba.pydata.org/numba-doc/0.12.2/index.html)  Example 1  Plain Python version<jupyter_code>def matrix_multiply(A, B):
    m, n = A.shape
    n, p = B.shape
    C = np.zeros((m, p))
    for i in range(m):
        for j in range(p):
            for k in range(n):
                C[i,j] += A[i,k] * B[k, j]
    return C
A = np.random.random((30, 50))
B = np.random.random((50, 40))<jupyter_output><empty_output><jupyter_text>Numba jit version<jupyter_code>import numba
from numba import jit
@jit
def matrix_multiply_numba(A, B):
    m, n = A.shape
    n, p = B.shape
    C = np.zeros((m, p))
    for i in range(m):
        for j in range(p):
            for k in range(n):
                C[i,j] += A[i,k] * B[k, j]
    return C<jupyter_output><empty_output><jupyter_text>We can remove the cost of indexing a matrix in the inner loop<jupyter_code>@jit
def matrix_multiply_numba2(A, B):
    m, n = A.shape
    n, p = B.shape
    C = np.zeros((m, p))
    for i in range(m):
        for j in range(p):
            d = 0.0
            for k in range(n):
                d += A[i,k] * B[k, j]
            C[i,j] = d
    return C
%timeit -r3 -n3 matrix_multiply(A, B)
%timeit -r3 -n3 matrix_multiply_numba(A, B)
%timeit -r3 -n3 matrix_multiply_numba2(A, B)<jupyter_output>55.7 ms ± 399 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)
The slowest run took 1715.66 times longer than the fastest. This could mean that an intermediate result is being cached.
33.6 ms ± 47.4 ms per loop (mean ± std. dev. of 3 runs, 3 loops each)
The slowest run took 1230.16 times longer than the fastest. This could mean that an intermediate result is being cached.
20 ms ± 28.2 ms per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Numpy version<jupyter_code>def matrix_multiply_numpy(A, B):
    return A.dot(B)<jupyter_output><empty_output><jupyter_text>Check that outputs are the same<jupyter_code>assert_almost_equal(matrix_multiply(A, B), matrix_multiply_numba(A, B))
assert_almost_equal(matrix_multiply(A, B), matrix_multiply_numpy(A, B))
%timeit -r3 -n3 matrix_multiply_numba(A, B)
report([matrix_multiply, matrix_multiply_numba, matrix_multiply_numba2, matrix_multiply_numpy], A, B)<jupyter_output>matrix_multiply: 1.0
matrix_multiply_numba: 850.3
matrix_multiply_numba2: 1025.4
matrix_multiply_numpy: 979.2<jupyter_text>Pre-compilation by giving specific signature<jupyter_code>@jit('double[:,:](double[:,:], double[:,:])')
def matrix_multiply_numba_1(A, B):
    m, n = A.shape
    n, p = B.shape
    C = np.zeros((m, p))
    for i in range(m):
        for j in range(p):
            d = 0.0
            for k in range(n):
                d += A[i,k] * B[k, j]
            C[i,j] = d
    return C
%timeit -r3 -n3 matrix_multiply_numba2(A, B)
%timeit -r3 -n3 matrix_multiply_numba_1(A, B)<jupyter_output>51.4 µs ± 2.59 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)
72.3 µs ± 31 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Example 2: Using nopython  Vectorized Python version<jupyter_code>def mc_pi(n):
    x = np.random.uniform(-1, 1, (n,2))
    return 4*np.sum((x**2).sum(1) < 1)/n
n = int(1e6)
mc_pi(n)
%timeit mc_pi(n)<jupyter_output>55.1 ms ± 861 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)<jupyter_text>Numba on vectorized version<jupyter_code>@jit
def mc_pi_numba(n):
    x = np.random.uniform(-1, 1, (n,2))
    return 4*np.sum((x**2).sum(1) < 1)/n
%timeit mc_pi_numba(n)<jupyter_output>41.3 ms ± 1.04 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Using nopythonUsing nopython, either with the `@njit` decorator or with `@jit(nopython = True)`, tells `numba` to not use any Python objects in the C code, but only native C types. If `numba` cannot do this, it will raise an error. It is usually useful to run this, so you are aware of bottlenecks in your code.<jupyter_code>@jit(nopython=True)
def mc_pi_numba_njit(n):
    x = np.random.uniform(-1, 1, (n,2))
    return 4*np.sum((x**2).sum(1) < 1)/n
%timeit mc_pi_numba_njit(n)<jupyter_output>41.9 ms ± 1.91 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Numba on unrolled version numpa works even better on code that looks like C (explicit loops for everything)<jupyter_code>@jit(nopython=True)
def mc_pi_numba_unrolled(n):
    s = 0
    for i in range(n):
        x = np.random.uniform(-1, 1)
        y = np.random.uniform(-1, 1)
        if (x*x + y*y) < 1:
            s += 1
    return 4*s/n
mc_pi_numba_unrolled(n)
%timeit -r3 -n3 mc_pi_numba_unrolled(n)<jupyter_output>14.5 ms ± 105 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Usig cache=TrueThis stores the compiled function in a file and avoids re-compilation on re-running a Python program.<jupyter_code>@jit(nopython=True, cache=True)
def mc_pi_numba_unrolled_cache(n):
    s = 0
    for i in range(n):
        x = np.random.uniform(-1, 1)
        y = np.random.uniform(-1, 1)
        if (x*x + y*y) < 1:
            s += 1
    return 4*s/n
%timeit -r3 -n3 mc_pi_numba_unrolled_cache(n)<jupyter_output>15.5 ms ± 1.51 ms per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Simple parallel loops with `numba`<jupyter_code>from numba import njit, prange
@njit()
def sum_rows_range(A):
    s = 0
    for i in range(A.shape[0]):
        s += np.sum(np.exp(np.log(np.sqrt(A[i]**2.0))))
    return s
@njit(parallel=True)
def sum_rows_prange(A):
    s = 0
    for i in prange(A.shape[0]):
        s += np.sum(np.exp(np.log(np.sqrt(A[i]**2.0))))
    return s
A = np.random.randint(0, 10, (800, 100000))
A.shape<jupyter_output><empty_output><jupyter_text>Run once so that compile times excluded in benchmarking<jupyter_code>sum_rows_range(A), sum_rows_prange(A)
%%time

sum_rows_range(A)
%%time

sum_rows_prange(A)<jupyter_output>CPU times: user 5.28 s, sys: 80 ms, total: 5.36 s
Wall time: 662 ms<jupyter_text>Using numba vectorize and guvectoize----Sometimes it is convenient to use `numba` to convert functions to vectorized functions for use in `numpy`. See [documentation](http://numba.pydata.org/numba-doc/dev/user/vectorize.html) for details.<jupyter_code>from numba import int32, int64, float32, float64<jupyter_output><empty_output><jupyter_text>Using `vectorize`<jupyter_code>@numba.vectorize()
def f(x, y):
    return np.sqrt(x**2 + y**2)
xs = np.random.random(10)
ys = np.random.random(10)
np.array([np.sqrt(x**2 + y**2) for (x, y) in zip(xs, ys)])
f(xs, ys)<jupyter_output><empty_output><jupyter_text>Adding function signatures<jupyter_code>@numba.vectorize([float64(float64, float64),
                  float32(float32, float32),
                  float64(int64, int64),
                  float32(int32, int32)])
def f_sig(x, y):
    return np.sqrt(x**2 + y**2)
f_sig(xs, ys)<jupyter_output><empty_output><jupyter_text>Using `guvectorize`  **Create our own version of inner1d**Suppose we have two matrices, each with `m` rows. We may want to calculate an "row-wise" inner product, that is, generate a scalar for each pair of row vectors. We cannot use `@vectorize` because the elements are not scalars.The *layout* `(n),(n)->()` says the function to be vectorized takes two `n`-element one dimensional arrays  `(n)` and returns a scalar `()`. The type *signature* is a list that matches the order of the *layout*.<jupyter_code>@numba.guvectorize([(float64[:], float64[:], float64[:])], '(n),(n)->()')
def nb_inner1d(u, v, res):
    res[0] = 0
    for i in range(len(u)):
        res[0] += u[i]*v[i] # Pass in 3 arguments, but last one is the return value
xs = np.random.random((3,4))
nb_inner1d(xs, xs)<jupyter_output><empty_output><jupyter_text>**Check**<jupyter_code>from numpy.core.umath_tests import inner1d
inner1d(xs,xs)<jupyter_output><empty_output><jupyter_text>Alternative to deprecated `inner1d` using Einstein summation notationFor more on how to use Einstein notation, see the help documentation and [here](https://rockt.github.io/2018/04/30/einsum)<jupyter_code>np.einsum('ij,ij->i', xs, xs)
%timeit -r3 -n3 nb_inner1d(xs, xs)
%timeit -r3 -n3 inner1d(xs, xs)<jupyter_output>The slowest run took 4.13 times longer than the fastest. This could mean that an intermediate result is being cached.
3.11 µs ± 2.18 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>**Create our own version of matrix_multiply**<jupyter_code>@numba.guvectorize([(int64[:,:], int64[:,:], int64[:,:])], 
                    '(m,n),(n,p)->(m,p)')
def nb_matrix_multiply(u, v, res):
    m, n = u.shape
    n, p = v.shape
    for i in range(m):
        for j in range(p):
            res[i,j] = 0
            for k in range(n):
                res[i,j] += u[i,k] * v[k,j]
xs = np.random.randint(0, 10, (5, 2, 3))
ys = np.random.randint(0, 10, (5, 3, 2))
nb_matrix_multiply(xs, ys)<jupyter_output><empty_output><jupyter_text>**Check**<jupyter_code>from numpy.core.umath_tests import matrix_multiply
matrix_multiply(xs, ys)
%timeit -r3 -n3 nb_matrix_multiply(xs, ys)
%timeit -r3 -n3 matrix_multiply(xs, ys)<jupyter_output>3.62 µs ± 2.32 µs per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Parallelization with vectorize and guvectorizeIf you have an NVidia graphics card and CUDA drivers installed, you can also use `target = 'cuda'`.<jupyter_code>@numba.vectorize([float64(float64, float64),
                  float32(float32, float32),
                  float64(int64, int64),
                  float32(int32, int32)],
                 target='parallel')
def f_parallel(x, y):
    return np.sqrt(x**2 + y**2)
xs = np.random.random(int(1e8))
ys = np.random.random(int(1e8))
%timeit -r3 -n3 f(xs, ys)
%timeit -r3 -n3 f_parallel(xs, ys)<jupyter_output>292 ms ± 6.36 ms per loop (mean ± std. dev. of 3 runs, 3 loops each)<jupyter_text>Mandelbrot example with `numba` **Pure Python**<jupyter_code># color function for point at (x, y)
def mandel(x, y, max_iters):
    c = complex(x, y)
    z = 0.0j
    for i in range(max_iters):
        z = z*z + c
        if z.real*z.real + z.imag*z.imag >= 4:
            return i
    return max_iters
def create_fractal(xmin, xmax, ymin, ymax, image, iters):
    height, width = image.shape
    
    pixel_size_x = (xmax - xmin)/width
    pixel_size_y = (ymax - ymin)/height
        
    for x in range(width):
        real = xmin + x*pixel_size_x
        for y in range(height):
            imag = ymin + y*pixel_size_y
            color = mandel(real, imag, iters)
            image[y, x]  = color    
gimage = np.zeros((1024, 1536), dtype=np.uint8)
xmin, xmax, ymin, ymax = np.array([-2.0, 1.0, -1.0, 1.0]).astype('float32')
iters = 50

start = time.time()
create_fractal(xmin, xmax, ymin, ymax, gimage, iters)
dt = time.time() - start

print("Mandelbrot created on CPU in %f s" % dt)
plt.grid(False)
plt.imshow(gimage, cmap='jet')
pass<jupyter_output>Mandelbrot created on CPU in 15.755038 s<jupyter_text>**Numba**<jupyter_code>from numba import uint32, float32<jupyter_output><empty_output><jupyter_text>**The jit decorator can also be called as a regular function**<jupyter_code>mandel_numba = jit(uint32(float32, float32, uint32))(mandel)
@jit
def create_fractal_numba(xmin, xmax, ymin, ymax, image, iters):
    height, width = image.shape
    
    pixel_size_x = (xmax - xmin)/width
    pixel_size_y = (ymax - ymin)/height
        
    for x in range(width):
        real = xmin + x*pixel_size_x
        for y in range(height):
            imag = ymin + y*pixel_size_y
            color = mandel_numba(real, imag, iters)
            image[y, x]  = color  
gimage = np.zeros((1024, 1536), dtype=np.uint8)
xmin, xmax, ymin, ymax = np.array([-2.0, 1.0, -1.0, 1.0]).astype('float32')
iters = 50

start = time.time()
create_fractal_numba(xmin, xmax, ymin, ymax, gimage, iters)
dt = time.time() - start

print("Mandelbrot created wiht Numba in %f s" % dt)
plt.grid(False)
plt.imshow(gimage, cmap='jet')
pass<jupyter_output>Mandelbrot created wiht Numba in 0.281742 s
{'fraction_non_alphanumeric': 0.1345668629100084, 'fraction_numerical': 0.03336136809643959, 'mean_word_length': 3.7754350736278446, 'pattern_counts': {'":': 0, '<': 103, '<?xml version=': 0, '>': 103, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '28', 'dir': 'jupyter-structured-clean-dedup', 'id': '31523', 'license': 'MIT', 'n_tokens_mistral': 5753, 'n_tokens_neox': 5471, 'n_words': 1560, 'path': 'notebooks/copies/lectures/S08B_Numba.ipynb', 'repo_name': 'robkravec/sta-663-2021', 'size': '14268'}