高效遍历Scipy稀疏CSR矩阵中每一行的非零元素

Question

我正在处理一个以CSR（压缩稀疏行）格式存储的scipy.sparse矩阵，并且需要逐行迭代以获取非零元素的索引及其对应的值。

虽然我已经了解到有方法可以有效地遍历CSR矩阵（参见：在Stack Overflow上的讨论），但目前尚未成功调整这些方法来满足我的需求。

现在，我是这样做的：

# 假设'matrix'是一个scipy.sparse.csr_matrix对象
for index in tqdm(range(matrix.shape[0]), desc="更新值", leave=False):
    # 获取当前行
    row = matrix.getrow(index)
    # 获取该行中非零元素的索引
    values_indices = row.indices
    # 对索引和值进行进一步处理...

请问有没有更高效的方式来实现这一目标？

phd · Answer

我发现了一种比最初使用的方法更为高效的遍历scipy稀疏CSR矩阵每一行中非零元素的方法。以下是两种方法及其性能对比摘要：

初始方法（使用getrow）：

# csr_matrix 是一个 scipy.sparse.csr_matrix 对象
for user_index in range(csr_matrix.shape[0]):
    row = csr_matrix.getrow(user_index)
    values_indices = row.indices

这种方法对于一个具有10,000行和5,000列的矩阵耗时约2.9309秒。

第二种基于这个链接的方法：

start_time = time.time()
old_i = None
indices = []
values = []

for i, j, v in zip(coo_matrix.row, coo_matrix.col, coo_matrix.data):
    if i != old_i:
        if old_i is not None:
            other_method(indices, values)
        indices = [j]
        values = [v]
    else:
        indices.append(j)
        values.append(v)
    old_i = i

# 处理最后一组数据
if indices and values:
    other_method(indices, values)

end_time = time.time()

这种方法对于同样大小的矩阵仅耗时0.0030秒。

然而，我不确定这两种方法是否都是最佳实践。

Rich Scriven · Answer

能够直接从CSR格式矩阵中获取每行的内容而无需转换为COO格式并遍历矩阵中的值，这样做通常更快，原因如下：

首先，将CSR格式转换为COO格式会消耗额外的时间。
其次，在COO循环中，为了区分不同行的数据，必须比较每一行的行坐标以确定它是否与前一行不同。相反，CSR格式本身就包含了每个行起始和结束位置的信息。
第三，通过切片操作可以从数组中获取代表某行数据的视图，而无需实际复制任何数据。

以下函数利用indptr属性直接找到每行的起始和结束位置，并对每行的值和索引调用指定的函数：

def get_matrix_rows(matrix, func):
    rows = matrix.shape[0]
    for index in range(rows):
        indptr_start = matrix.indptr[index]
        indptr_end = matrix.indptr[index + 1]
        values = matrix.data[indptr_start:indptr_end]
        indices = matrix.indices[indptr_start:indptr_end]
        func(indices, values)

需要注意的是，如果矩阵某行没有元素，此实现会传入空的indices和values数组调用函数；而在另一种实现中，则会跳过所有不含非零值的行。

下面是一段基准测试脚本，用于比较两种实现方式。其假设条件包括：

矩阵是10,000行 x 5,000列的CSR格式矩阵，由随机生成的1%密度的非零值组成。
在计算COO方法的时间时，包含了将矩阵转换为COO格式的过程。
每个函数都需找出各行的值和列索引，并调用给定的函数。

代码如下：

import scipy.sparse
matrix = scipy.sparse.random(10000, 5000, format='csr', density=0.01)

# 示例无操作函数
def donothing(*args):
    pass

# 初始方法，使用getrow()
def get_matrix_original(matrix, func):
    for index in range(matrix.shape[0]):
        row = matrix.getrow(index)
        indices = row.indices
        values = row.data
        func(indices, values)

# COO格式迭代方法
def get_matrix_rows_coo(matrix, func):
    coo_matrix = matrix.tocoo()
    old_i = None
    indices = []
    values = []

    for i, j, v in zip(coo_matrix.row, coo_matrix.col, coo_matrix.data):
        if i != old_i:
            if old_i is not None:
                func(indices, values)
            indices = [j]
            values = [v]
        else:
            indices.append(j)
            values.append(v)
        old_i = i

    # 处理最后一组数据
    if indices and values:
        func(indices, values)

# 直接读取CSR格式方法
def get_matrix_rows_njo(matrix, func):
    rows = matrix.shape[0]
    for index in range(rows):
        indptr_start = matrix.indptr[index]
        indptr_end = matrix.indptr[index + 1]
        values = matrix.data[indptr_start:indptr_end]
        indices = matrix.indices[indptr_start:indptr_end]
        func(indices, values)

# 执行时间测试
%timeit get_matrix_original(matrix, donothing)
%timeit get_matrix_rows_coo(matrix, donothing)
%timeit get_matrix_rows_njo(matrix, donothing)

基准测试结果显示：

.getrow() 方法：634毫秒±16.8毫秒/次循环（平均值±标准差，7次运行，每次1个循环）
COO格式并迭代的方法：270毫秒±4.4毫秒/次循环（平均值±标准差，7次运行，每次1个循环）
CSR格式直接读取的方法：12.4毫秒±112微秒/次循环（平均值±标准差，7次运行，100个循环）

特别注意，在非常低的密度条件下（大约₀.05%的非零值），由于COO方法不需要对空行做任何处理，所以其速度可能超过CSR方法。