Numba，Numpy：在并行 @jit 函数中调用 @guvectorized 函数答案

【问题标题】：Numba, Numpy: Call @guvectorized function in a parallel @jit functionNumba，Numpy：在并行 @jit 函数中调用 @guvectorized 函数
【发布时间】：2021-01-28 21:28:28
【问题描述】：

考虑以下计算布尔向量之间相似度的简单函数。

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

@guvectorize([(boolean[:], boolean[:], float64[:])], '(n),(n)->()')
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

我们可以用两个布尔向量调用这个函数，也可以用一个布尔向量和一个布尔向量数组来计算几个向量的相似度（在老式的 numpy 风格中）：

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

tanimoto(fp1, [fp2, fp3])

现在，我正在计算相似度矩阵的上三角形，该矩阵基本上是布尔向量数组的所有成对相似度。

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

similarityMatrix(np.array([fp1,fp2,fp3]))

但是，我似乎无法正确地 @jit-compile 和并行化 similarityMatrix 函数。当我添加 @jit(parallel=True) 注释时，我收到以下几条消息：

编译正在回退到启用循环提升的对象模式，因为函数“similarityMatrix”由于以下原因导致类型推断失败：无类型全局名称“tanimoto”：无法确定的 Numba 类型>

问题：有没有办法让内部 prange-loop 并行运行？还有其他我不知道的事情可以让这尽可能快吗？

旁注：我知道 scipy 的 pdist 函数。这里的目标是更好地了解 Numba 并深入了解我目前缺少的内容。

【问题讨论】：

标签： python numpy numba

【解决方案1】：

回答我自己的问题。在 Numba Gitter 聊天中，有人指出：

问题是 guvectorize 正在生成一个真正的 NumPy ufunc 实例，这不是 Numba 可以理解的。

他们建议尝试将 guvectorized 函数实现为一个简单的循环。令人惊讶的是，这确实提供了更好的性能。计算大约 7k 个值的相似度矩阵现在只需要 13 秒，而矢量化版本需要 1:19 分钟。计算调用相似性度量 26,263,128 次，我对性能非常满意。

布尔向量的 Tanimoto 相似度矩阵的最终实现现在如下所示：

@njit
def tanimotoSimilarity(fp: np.ndarray, fps: np.ndarray) -> np.ndarray:
    result = np.zeros(fps.shape[0])
    for i in prange(0, result.shape[0]):
        den = np.sum(np.bitwise_or(fp, fps[i]))
        if den != 0.0:
            result[i] = np.sum(np.bitwise_and(fp, fps[i])) / den
        else:
            result[i] = 0.0
    return result


@njit(parallel=True, cache=True)
def similarityMatrix(fp_list: np.ndarray) -> np.ndarray:
    m = fp_list.shape[0]
    sm = 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):
        sm[idx[i]:idx[i+1]] = tanimotoSimilarity(fp_list[i], fp_list[i + 1:])
    return sm

【讨论】：