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

Question

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

我创建ROW_INDEX矢量的第一种方法是：

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

毫不奇怪，这相当缓慢。然后我尝试了旧的速度交换空间：

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)