Ускорить левый матмуль для scipy.sparse.csr_matrix

Question

Мне нужно выполнить следующее матричное умножение: x * A[idx], где A - это scipy.sparse.csr_matrix, а idx - это np.array индекс. Я не могу изменить его на csc_matrix из-за индексации. Это кажется на 50 медленнее, чем умножение правой матрицы A[idx] * x, и лишь немного быстрее, чем умножение левой матрицы (по полному индексу) 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