alko/seberg 的方法略有不同。我发现 for 循环令人不安,所以我花了今天早上的大部分时间想办法摆脱它。以下方法并不总是比其他方法快。归零的行越多,矩阵越稀疏,它的性能就越好:
def csr_zero_rows(csr, rows_to_zero):
rows, cols = csr.shape
mask = np.ones((rows,), dtype=np.bool)
mask[rows_to_zero] = False
nnz_per_row = np.diff(csr.indptr)
mask = np.repeat(mask, nnz_per_row)
nnz_per_row[rows_to_zero] = 0
csr.data = csr.data[mask]
csr.indices = csr.indices[mask]
csr.indptr[1:] = np.cumsum(nnz_per_row)
并测试两种方法:
rows, cols = 334863, 334863
a = sps.rand(rows, cols, density=0.00001, format='csr')
b = a.copy()
rows_to_zero = np.random.choice(np.arange(rows), size=10000, replace=False)
In [117]: a
Out[117]:
<334863x334863 sparse matrix of type '<type 'numpy.float64'>'
with 1121332 stored elements in Compressed Sparse Row format>
In [118]: %timeit -n1 -r1 csr_rows_set_nz_to_val(a, rows_to_zero)
1 loops, best of 1: 75.8 ms per loop
In [119]: %timeit -n1 -r1 csr_zero_rows(b, rows_to_zero)
1 loops, best of 1: 32.5 ms per loop
而且当然:
np.allclose(a.data, b.data)
Out[122]: True
np.allclose(a.indices, b.indices)
Out[123]: True
np.allclose(a.indptr, b.indptr)
Out[124]: True