sparse adjacency matrix from distance metric and thresholding in numba

calculate_adjacency_matrix.py

from numba import njit, jit, prange
import numpy as np
from numba.pycc import CC
from scipy.sparse import lil_matrix

cc = CC('adjacency_utils')

@cc.export('calc_dist', 'f8(f8[:], f8[:])')
@jit("f8(f8[:], f8[:])")
def calc_dist(u, v):
    '''Euclidean distance (without sqrt)

    Example:
    --------
    >> calc_dist(X[0, :], X[1, :])
    12.795783809844064

    '''
    d = u - v
    return np.sum(d * d)

@jit(nopython=False, parallel=True, forceobj=True)
def calculate_adjacency(X, threshold=0.5):
    '''Calculate an adjacency matrix
    given a feature matrix
    '''
    n_rows = X.shape[0]
    A = lil_matrix((n_rows, n_rows), dtype=np.int8)
    for i in prange(n_rows):
        for i2 in range(i+1, n_rows):
            d = calc_dist(X[i, :], X[i2, :])
            if d < threshold:
                A[i, i2] = 1.0
                A[i2, i] = 1.0
    return A

cc.compile()