similarity.py

from numba import float64, boolean, prange, guvectorize, njit, set_num_threads
import numpy as np


@guvectorize([(boolean[:], boolean[:], float64[:])], '(n),(n)->()', nopython=True)
def tanimoto(fp_1, fp_2, res):
    bw_or = np.sum(np.bitwise_or(fp_1, fp_2))
    if bw_or != 0.0:
        res[0] = np.sum(np.bitwise_and(fp_1, fp_2)) / bw_or
    else:
        res[0] = 0.0


@njit(parallel=True, cache=True)
def similarityMatrix(fp_list: np.ndarray) -> np.ndarray:
    m = fp_list.shape[0]
    dm = np.zeros((m * (m - 1)) // 2)
    idx = [int((2*i*m-i*i-i)/2) for i in range(0, m)]
    for i in prange(0, m - 1):
        dm[idx[i]:idx[i+1]] = tanimoto(fp_list[i], fp_list[i+1:])
    return dm


fp1 = np.array([True, False, False, True, False])
fp2 = np.array([False, False, False, True, False])
fp3 = np.array([True, False, True, True, False])
similarityMatrix(np.asarray([fp1, fp2, fp3]))