Numba中的稀疏矩阵

Question

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

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

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

Answer 1

如果您所要做的只是迭代 CSR 矩阵的值，则可以将属性 data、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)

Answer 2

您可以将稀疏矩阵的数据作为纯 numpy 或 python 访问。例如

M=sparse.csr_matrix([[1,0,0],[1,0,1],[1,1,1]])
ML = M.tolil()

for d,r in enumerate(zip(ML.data,ML.rows))
    # d,r are lists
    dr = np.array([d,r])
    print dr

产生：

[[1]
 [0]]
[[1 1]
 [0 2]]
[[1 1 1]
 [0 1 2]]

当然，numba 可以处理使用这些数组的代码，当然，前提是它不期望每行具有相同大小的数组。

该lil格式存储值 2 个对象 dtype 数组，其中数据和索引按行存储列表。