Compute difference of dense matrix at non-zero elements of another sparse matrix

Question

I have a n x n matrix W a n-diemensional vector U and a scalar p > 1. I want to compute W * (np.abs(U[:,np.newaxis] - U) ** p), where *, ** and abs are understood element wise (as in numpy).

Now the problem is that U[:,np.newaxis] - U does not fit into memory. However, W is a sparse matrix (scipy.sparse), so I don't actually have to compute all entries of U[:,np.newaxis] - U, but only those, where W is not zero.

How can I compute W * (np.abs(U[:,np.newaxis] - U) ** p) most efficiently in terms of computation time and memory, ideally doing only sparse operations, without a step through numpy?

Answer 1

To make use of the sparseness you could apply the distributive law so

 C = W * (np.abs(U[:,np.newaxis] - U) ** p)

would result in

rtW = W**(1/p)
C = (rtW * np.abs(U[:,np.newaxis] - rtW * U) ** p

note that this only makes sense as long the entries of W aren't negative. But we can remedy that by using

rtW = np.abs(W)**(1/p)
signW = np.sign(W)
C = signW * (rtW * np.abs(U[:,np.newaxis] - rtW * U) ** p