如何在 python 中高效地创建 CSR row_index 向量？

Question

我正在尝试创建一个 CSR matrix，其中 m 行，n 列，填充零和一（每列最多一个）。我有一个 numpy 数组 idcs，其中包含我的 1 所在的索引，范围从 x 到 n。

我创建 ROW_INDEX 向量的第一种方法类似于：

ROW_INDEX=np.zeros(m+1)
for i in idcs: ROW_INDEX[i+1:]+=1

不出所料，这相当慢。然后我尝试了旧的 space 进行速度交换：

ROW_INDEX=np.fromfunction(lambda i,j: i>idcs[j],(m+1,n),dtype='uintc')
ROW_INDEX=np.sum(ROW_INDEX,1)

但是，m 和 n 都是 10^5，所以上面的代码引发了一个 MemoryError - 尽管大矩阵在技术上只是布尔值。

我觉得我在这里错过了一些明显的东西。任何人都有更聪明的解决方案，还是我应该增加内存？

最终目的是创建一个 PETSc.Mat，希望是并行的，从类似 B=petsc4py.Mat().createAIJ([m, n],csr=[ROW_INDEX,COL_INDEX,V]) 的东西开始。我发现关于这个主题的文档很少，也欢迎在这方面提供任何帮助。

Answer 1

我想你正在寻找这样的东西？

ROW_INDEX=np.zeros(m+1)
np.add.at(ROW_INDEX, idcs+1, 1)
np.cumsum(ROW_INDEX, out=ROW_INDEX)

How to efficiently create a CSR row_index vector in python?