SciPy CSR 稀疏矩阵高效行迭代：直接利用内部结构优化非零元素访问

import scipy.sparse
from tqdm import tqdm

# 'matrix' is a 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
#     # Further processing...

# coo_matrix = matrix.tocoo()
# for i, j, v in zip(coo_matrix.row, coo_matrix.col, coo_matrix.data):
#     # 需要手动追踪行边界
#     pass

def get_matrix_rows_optimized(matrix, func):
    """
    高效迭代 CSR 矩阵的每一行非零元素。
    直接利用 CSR 矩阵的 indptr, data, indices 属性。

    Args:
        matrix (scipy.sparse.csr_matrix): 要迭代的 CSR 矩阵。
        func (callable): 对每行非零元素的索引和值进行操作的函数。
                         函数签名应为 func(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. 创建一个稀疏 CSR 矩阵
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_optimized(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(".getrow() method:")
print(timeit.timeit("get_matrix_original(matrix, donothing)", globals=globals(), number=1)) # number=1 for larger ops

print("\nCOO and iterate method:")
print(timeit.timeit("get_matrix_rows_coo(matrix, donothing)", globals=globals(), number=1))

print("\nCSR optimized method:")
print(timeit.timeit("get_matrix_rows_optimized(matrix, donothing)", globals=globals(), number=100)) # number=100 for faster ops

.getrow() method
0.634 seconds

COO and iterate method
0.270 seconds

CSR optimized method
0.012 seconds (for 100 loops, so ~0.00012 seconds per loop)