在没有 scipy.sparse 的情况下向量化稀疏和

Vectorize sparse sum without scipy.sparse

我正在尝试进行空间导数并且几乎成功地从我的代码中获取所有循环,但是当我尝试在最后总结所有内容时我遇到了问题。

我有一组 N~=250k 个节点。我找到了在一定搜索距离内具有 i.size=j.size=~7.5M 的节点对的索引 i,j,最初来自 np.triu_indices(n,1) 并通过一系列布尔掩码以清除互不影响的节点.现在我想总结一下其他节点对每个节点的影响。

我目前有这个:

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)])

这很慢。我想要的是矢量化的东西。如果我有 scipy 我可以做到

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)

但我是在不包含 scipy 的 Abaqus 实现中完成这一切的。有什么办法可以只用 numpy 吗?

方法 #1 : 这是一种使用 matrix-multiplication and broadcasting -

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

方法 #2 : 对于列数较少的情况,我们可以使用 np.bincount 对每一列进行这样的 bin-based 求和,就像这样-

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

这不是 turn-key 解决方案,而是添加稀疏矩阵列的解决方案。它本质上是计算和利用 csc 表示

def sparse_col_sums(i, j, a, N):
    order = np.lexsort(j, i)
    io, jo, ao = i[order], j[order], a[order]
    col_bnds = io.searchsorted(np.arange(N))
    return np.add.reduceat(ao, col_bnds)