问题描述
我需要从循环中的稀疏矩阵表达式 x
中求解 A*x = B
,其中 A
是 Scipy CSC 稀疏矩阵,B
是 Numpy 一维数组。 A
和 B
都是大约 50 万行。基本上,我需要在每个循环中更新 B
。因此更新 B
的速度至关重要。现在,我的方法是在每个循环中定义 csc_matrix
,然后将其转换为如下所示的 1D Numpy 数组,这在时间方面非常昂贵:
B = csc_matrix((data,(row,col)),shape=(500000,1),dtype='complex128').toarray()[:,0];
请注意:
-
row
有很多重复的索引,比如[0,1,2,3,3....],
-
col
是[0,.......0];
解决方法
假设 col
只包含零,data
/row
/col
是 Numpy 数组,您希望 B
存储为 Numpy 数组。您可以使用 Numba 有效地生成 B
。方法如下:
import numba
# Works in-place to avoid any slow allocation in the critical loop.
# Note that the type of row may be different.
@nb.njit(void(nb.complex128[:],nb.complex128[:],nb.int64[:]))
def updateVector(B,data,row):
B.fill(0.)
for i in range(len(row)):
B[row[i]] += data[i]
updateVector
就地更新 B
的值。这假设 B
之前已以正确的大小分配(例如使用 B = np.empty(500000,dtype=np.complex128)
)。
在我的机器上,使用以下配置可以 快 14 倍:
row = np.random.randint(0,500000,size=100000)
col = np.zeros(100000,dtype=np.int64)
data = np.random.rand(100000) + np.random.rand(100000) * 1j