Mean of non zero values in sparse matrix?

Question

I'm trying to calculate the mean of non-zero values in each row of a sparse row matrix. Using the matrix's mean method doesn't do it:

>>> from scipy.sparse import csr_matrix
>>> a = csr_matrix([[0, 0, 2], [1, 3, 8]])
>>> a.mean(axis=1)
matrix([[ 0.66666667],
        [ 4.        ]])

The following works but is slow for large matrices:

>>> import numpy as np
>>> b = np.zeros(a.shape[0])
>>> for i in range(a.shape[0]):
...    b[i] = a.getrow(i).data.mean()
... 
>>> b
array([ 2.,  4.])

Could anyone please tell me if there is a faster method?

Answer 1

With a CSR format matrix, you can do this even more easily:

sums = a.sum(axis=1).A1
counts = np.diff(a.indptr)
averages = sums / counts

Row-sums are directly supported, and the structure of the CSR format means that the difference between successive values in the indptr array correspond exactly to the number of nonzero elements in each row.