使用稀疏scipy矩阵的广播

Question

I have a numpy array Z with shape (k,N) and a second array X with shape (N,n) .

Using numpy broadcasting, I can easily obtain a new array H with shape (n,k,N) whose slices are the array Z whose rows have been multiplied by the columns of X :

H = Z.reshape((1, k, N)) * X.T.reshape((n, 1, N))

This works fine and is surprisingly fast. Now, X is extremely sparse, and I want to further speed up this operation using sparse matrix operations.

However if I perform the following operations:

import scipy.sparse as sprs
spX = sprs.csr_matrix(X)
H = (Z.reshape((1,k,N))*spX.T.reshape((n,1,N))).dot(Z.T)

I get the following error:

Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File "C:\Python27\lib\site-packages\scipy\sparse\base.py", line 126, in reshape
    self.__class__.__name__)
NotImplementedError: Reshaping not implemented for csc_matrix.

Is there a way to use broadcasting with sparse scipy matrices?

Answer 1

Scipy sparse matrices are limited to 2D shapes. But you can use Numpy in a "sparse" way:

H = np.zeros((n,k,N), np.result_type(Z, X))
I, J = np.nonzero(X)
Z_ = np.broadcast_to(Z, H.shape)
H[J,:,I] = Z_[J,:,I] * X[I,J,None]

Unfortunately the result H is still a dense array.

Nb indexing with None is a handy way to add a unit-length dimension at the desired axis. The order of the result when combining advanced indexing with slicing is explained in the docs .