Is there a numpy/scipy dot product for sparse matrix, calculating only the diagonal entries of the result?

Question

Imagine having 2 sparse matrix:

> A, A.shape = (n,m)
> B, B.shape = (m,n)

I would like to compute the dot product A*B, but then only keep the diagonal. The matrices being big, I actually don't want to compute other values than the ones in the diagonal.

This is a variant of the question Is there a numpy/scipy dot product, calculating only the diagonal entries of the result?

Where the most relevant answer seems to be to use np.einsum:

np.einsum('ij,ji->i', A, B)

However this does not work:

ValueError: einstein sum subscripts string contains too many subscripts for operand 0

The solution is to use todense(), but it increases a lot the memory usage: np.einsum('ij,ji->i', A.todense(), B.todense())

The other solution, that I currently use, is to iterate over all the rows of A and compute each product in the loop :

for i in range(len_A):
   result = np.float32(A[i].dot(B[:, i])[0, 0])
   ...

None of these solutions seems perfect. Is there an equivalent to np.einsum that could work with sparse matrices ?

Answer 1

[sum(A[i]*B.T[i]) for i in range(min(A.shape[0], B.shape[1]))]

otherwise this is faster:

l = min(A.shape[0], B.shape[1])
(A[np.arange(l)]*B.T[np.arange(l)]).sum(axis=1)

Answer 2

In general you shouldn't try to use numpy functions on the scipy.sparse arrays. In your case I'd first make sure both arrays actually have a compatible shape, that is

A, A.shape = (r,m)
B, B.shape = (m,r)

where r = min(n, p). Then we can compute the diagonal of the matrix product using

d = (A.multiply(B.T)).sum(axis=1)

Here we compute the entry wise row-column products, and manually sum them up. This avoids all the unnecessary computations you'd get using dot/@/*. (Note that unlike in numpy, both * and @ perform matrix multiplication.)