为 scipy.sparse.csr_matrix 加速左 matmul

Question

我需要执行以下矩阵乘法：x * A[idx] 其中 A 是 scipy.sparse.csr_matrix，idx 是 np.array 索引。由于索引，我无法将其更改为 csc_matrix。它似乎比右矩阵乘法 A[idx] * x 慢 50 倍，并且仅比左矩阵乘法（在整个索引上）u * A 稍快，即使在 len(idx) << A.shape[0] 时也是如此。我怎样才能加快速度？

Answer 1

由于我没有在其他地方找到解决方案，所以这就是我最终要做的。

使用numba，我写道：

@njit(nogil=True)
def fast_csr_vm(x, data, indptr, indices, d, idx):
    """
    Returns the vector matrix product x * M[idx]. M is described
    in the csr format.
    Returns x * M[idx]
    x: 1-d iterable
    data: data field of a scipy.sparse.csr_matrix
    indptr: indptr field of a scipy.sparse.csr_matrix
    indices: indices field of a scipy.sparse.csr_matrix
    d: output dimension
    idx: 1-d iterable: index of the sparse.csr_matrix
    """
    res = np.zeros(d)
    assert x.shape[0] == len(idx)
    for k, i in np.ndenumerate(idx):
        for j in range(indptr[i], indptr[i+1]):
            j_idx = indices[j]
            res[j_idx] += x[k] * data[j]
    return res

为 scipy.sparse.csr_matrix 加速左 matmul

Speed up left matmul for scipy.sparse.csr_matrix

python

scipy

sparse-matrix