加快scipy.sparse.csr_matrix的左matmul

Question

我需要执行以下矩阵乘法：x * A[idx]其中A是scipy.sparse.csr_matrix，而idx是np.array索引。由于建立索引，我无法将其更改为csc_matrix。即使A[idx] * u，它似乎也比右矩阵乘法len(idx) << A.shape[1]慢至少50倍。如何加快速度？

Answer 1

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

使用numba，我写道：

from numba import njit
@njit
def csr_left_mul(x, data, indptr, indices, d, idx):
    """
    :param x: numpy.ndarray
    :param A: scipy.sparse.csr_matrix
    :return: x.dot(A)
    """
    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