<jupyter_start><jupyter_text>Using `numba.jit` to speedup the computation of the Cityblock distance matrix In this notebook we implement a function to compute the Cityblock distance matrix using Numba's *just-it-time* compilation decorator. We compare it's performance to that of corresponding non-decorated NumPy function.We will use two Numba functions here. The decorator ` @numba.jit` and `numba.prange`.<jupyter_code>import numpy as np
import numba
def cityblock_python(x, y):
    """Naive python implementation."""

    num_samples, num_feat = x.shape
    dist_matrix = np.empty((num_samples, num_samples))
    for i in range(num_samples):
        for j in range(num_samples):
            r = 0.0
            for k in range(num_feat):
                r += np.abs(x[i][k] - y[j][k])
            dist_matrix[i][j] = r

    return dist_matrix


@numba.jit(nopython=True)
def cityblock_numba1(x, y):
    """Implementation with numba."""

    num_samples, num_feat = x.shape
    dist_matrix = np.empty((num_samples, num_samples))
    for i in range(num_samples):
        for j in range(num_samples):
            r = 0.0
            for k in numba.prange(num_feat):
                r += np.abs(x[i][k] - y[j][k])
            dist_matrix[i][j] = r

    return dist_matrix


@numba.jit(nopython=True)
def cityblock_numba2(x, y):
    """Implementation with numba and numpy."""

    num_samples, num_feat = x.shape
    dist_matrix = np.empty((num_samples, num_samples))
    for i in range(num_samples):
        for j in numba.prange(num_samples):
            dist_matrix[i][j] = np.linalg.norm(x[i] - y[j], 1)

    return dist_matrix<jupyter_output><empty_output><jupyter_text>NoteObserve that the inner loop, which is a reduction, is done with `numba.prange`. `numba.prange` automatically takes care of data privatization and reductions.<jupyter_code># Let's check that they all give the same result
a = 10. * np.random.random([100, 10])

print(np.abs(cityblock_python(a, a) - cityblock_numba1(a, a)).max())
print(np.abs(cityblock_python(a, a) - cityblock_numba2(a, a)).max())
nsamples = 200
nfeat = 25

x = 10. * np.random.random([nsamples, nfeat])

%timeit cityblock_python(x,x)
%timeit cityblock_numba1(x, x)
%timeit cityblock_numba2(x, x)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup