04_numba_functions.py

from numba.typed import List
from numba import njit

@njit
def boolean_index_numba(array):
    """
    Takes a numpy array and isolates all points that are within [0.2,0.4] for
    the first dimension and between [0.4,0.6] for the second dimension
    by creating a boolean index. 
    This function will be compiled with numba.
    """
    index = (array[:, 0] >= 0.2) & (array[:,  1] >= 0.4) & (array[:, 0] <= 0.4) & (array[:, 1] <= 0.6)

    return array[index]


@njit
def loop_numba(array):
    """
    Takes a numpy array and isolates all points that are within [0.2,0.4] for
    the first dimension and between [0.4,0.6] for the second dimension.
    This function will be compiled with numba.
    """
    filtered_list = List()
    for i in range(len(array)):
        if ((array[i][0] >= 0.2)
            and (array[i][1] >= 0.4)
            and (array[i][0] <= 0.4)
            and (array[i][1] <= 0.6)):
            filtered_list.append(array[i])
    return filtered_list


filtered_list = boolean_index_numba(array)
print('Boolean index with numba:\t', end='')
%timeit boolean_index_numba(array)

filtered_list = loop_numba(array)
print('Loop with numba:\t\t', end='')
%timeit loop_numba(array)