python - Numba 中的稀疏矩阵

Question

我希望使用 Numba ( http://numba.pydata.org/ ) 加速我的机器学习算法(用 Python 编写)。请注意，此算法将稀疏矩阵作为其输入数据。在我的纯 Python 实现中，我使用了来自 Scipy 的 csr_matrix 和相关类，但显然它与 Numba 的 JIT 编译器不兼容。

我还创建了自己的自定义类来实现稀疏矩阵(它基本上是(索引，值)对列表的列表)，但它再次与 Numba 不兼容(即，我收到一些奇怪的错误消息说它不识别扩展类型)

是否有一种仅使用与 Numba 兼容的 numpy(不求助于 SciPy)来实现稀疏矩阵的简单方法？任何示例代码将不胜感激。谢谢!

Answer 1

如果您只需迭代 CSR 矩阵的值，您可以将属性数据、indptr 和索引传递给函数而不是 CSR 矩阵对象。

from scipy import sparse
from numba import njit

@njit
def print_csr(A, iA, jA):
    for row in range(len(iA)-1):
        for i in range(iA[row], iA[row+1]):
            print(row, jA[i], A[i])

A = sparse.csr_matrix([[1, 2, 0], [0, 0, 3], [4, 0, 5]])
print_csr(A.data, A.indptr, A.indices)