Scipy Sparse CSR 矩阵非零元素行级遍历的最佳实践

import scipy.sparse
from tqdm import tqdm # 假设用于进度显示

# 'matrix' 是一个 scipy.sparse.csr_matrix
# for index in tqdm(range(matrix.shape[0]), desc="Updating values", leave=False):
#     row = matrix.getrow(index)
#     values_indices = row.indices
#     values = row.data
#     # 进一步处理...

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)

def iterate_csr_rows_efficiently(matrix, func):
    """
    高效遍历 CSR 矩阵的每一行非零元素。

    参数:
        matrix (scipy.sparse.csr_matrix): 要遍历的 CSR 稀疏矩阵。
        func (callable): 一个函数，接受两个参数 (indices, values)，
                         分别代表当前行的列索引和非零值。
    """
    rows = matrix.shape[0]
    for index in range(rows):
        # 使用 indptr 确定当前行在 data 和 indices 数组中的起始和结束位置
        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)

import scipy.sparse
import numpy as np
import timeit

# 1. 创建测试矩阵
matrix = scipy.sparse.random(10000, 5000, format='csr', density=0.01, random_state=42)

# 2. 定义一个空操作函数，用于基准测试
def donothing(*args):
    pass

# 3. 定义三种迭代方法

# 原始的 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_efficient_csr(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)

# 4. 执行基准测试
print("--- 性能基准测试结果 ---")

# 使用 timeit 模块进行测试
# 注意：在 Jupyter/IPython 环境中可以使用 %timeit，这里使用 timeit.timeit
num_runs = 7
num_loops_original = 1
num_loops_coo = 1
num_loops_efficient_csr = 100 # 显著更快，增加循环次数以获得更精确结果

time_original = timeit.timeit(lambda: get_matrix_original(matrix, donothing), number=num_loops_original)
print(f".getrow() 方法: {time_original / num_loops_original * 1000:.1f} ms ± ... (平均每次循环)")

time_coo = timeit.timeit(lambda: get_matrix_rows_coo(matrix, donothing), number=num_loops_coo)
print(f"COO 转换与迭代方法: {time_coo / num_loops_coo * 1000:.1f} ms ± ... (平均每次循环)")

time_efficient_csr = timeit.timeit(lambda: get_matrix_rows_efficient_csr(matrix, donothing), number=num_loops_efficient_csr)
print(f"CSR 直接属性访问方法: {time_efficient_csr / num_loops_efficient_csr * 1000:.1f} ms ± ... (平均每次循环)")

--- 性能基准测试结果 ---
.getrow() 方法: 634.0 ms ± ... (平均每次循环)
COO 转换与迭代方法: 270.0 ms ± ... (平均每次循环)
CSR 直接属性访问方法: 12.4 ms ± ... (平均每次循环)