Vectorize sparse sum without scipy.sparse

Question

I am trying to do spatial derivatives and almost managed to get all the loops out of my code, but when I try to sum everything up at the end I have a problem.

I have a set of N~=250k nodes. I have found indices i,j of node pairs with i.size=j.size=~7.5M that are within a certain search distance, originally coming from np.triu_indices(n,1) and passed through a series of boolean masks to wash out nodes not influencing each other. Now I want to sum up the influences on each node from the other nodes.

I currently have this:

def sparseSum(a,i,j,n):
    return np.array([np.sum(a[np.logical_or(i==k,j==k)],axis=0) for k in range(n)])

This is very slow. What I would like is something vectorized. If I had scipy I could do

def sparseSum(a,i,j,n):
    sp=scipy.sparse.csr_matrix((a,(i,j)),shape=(n,n))+ scipy.sparse.csr_matrix((a,(j,i)),shape=(n,n))
    return np.sum(sp, axis=0)

But I'm doing this all within an Abaqus implementation that doesn't include scipy. Is there any way to do this numpy-only?

Answer 1

Approach #1 : Here's an approach making use of matrix-multiplication and broadcasting -

K = np.arange(n)[:,None]
mask = (i == K) | (j == K)
out = np.dot(mask,a)

Approach #2 : For cases with a small number of columns, we can use np.bincount for such bin-based summing along each column, like so -

def sparseSum(a,i,j,n):
    if len(a.shape)==1:
        out=np.bincount(i,a,minlength=n)+np.bincount(j,a)
    else:
        ncols = a.shape[1]
        out = np.empty((n,ncols))
        for k in range(ncols):
            out[:,k] = np.bincount(i,a[:,k],minlength=n) + np.bincount(j,a[:,k])
    return out