Numba与CSR矩阵优化稀疏距离计算

优化稀疏成对距离计算

在数据分析和机器学习领域，我们经常需要计算两组向量集合 A 和 B 之间的成对距离。然而，在某些场景下，我们可能只对其中一小部分距离感兴趣，例如，通过一个稀疏掩码 M 来指定需要保留的距离对。如果直接计算所有可能的成对距离，然后通过掩码进行筛选，将导致大量的冗余计算和内存消耗，尤其当向量集合规模庞大时，这种低效性会变得尤为突出。

考虑以下示例，其中 A 和 B 是两组向量，M 是一个布尔掩码，指示哪些距离需要计算和保留：

import numpy as np

A = np.array([[1, 2], [2, 3], [3, 4]])                              # (3, 2)
B = np.array([[4, 5], [5, 6], [6, 7], [7, 8], [8, 9]])              # (5, 2)
M = np.array([[0, 0, 0, 1, 0], [1, 1, 0, 0, 0], [0, 0, 0, 0, 1]])   # (3, 5)

# 传统方法：计算所有距离
diff = A[:,None] - B[None,:]                                        # (3, 5, 2)
distances = np.linalg.norm(diff, ord=2, axis=2)                     # (3, 5)
masked_distances = distances * M                                    # (3, 5)

上述代码首先计算了 A 中每个向量与 B 中每个向量之间的所有差值，形成一个三维数组 diff。接着，对 diff 的每个“行-列”对计算其欧几里得范数，得到一个包含所有成对距离的二维矩阵 distances。最后，通过掩码 M 将不需要的距离置零。当 A 和 B 包含数千甚至数万个向量时，diff 和 distances 矩阵会变得非常巨大，导致内存溢出和计算时间过长。

另一种尝试是使用 np.vectorize 来条件性地计算差值：

def conditional_diff(a, b, m):
    if m:
        return a-b
    else:
        return np.zeros_like(a) # 确保返回相同形状的零向量

conditional_diff_vec = np.vectorize(conditional_diff, otypes=[float])
masked_diff = conditional_diff_vec(A[:,None], B[None,:], M[:,:,None])
masked_distances = np.linalg.norm(masked_diff, ord=2, axis=2)

尽管 np.vectorize 旨在将标量函数向量化，但其底层仍然涉及 Python 级别的循环和函数调用开销，对于大规模数据，性能甚至可能比直接的 NumPy 广播操作更差。因此，我们需要一种更根本的优化策略。

Numba 加速与 CSR 稀疏矩阵

解决上述问题的核心思想是：只计算那些由掩码 M 指定的、真正需要的距离，并将结果以内存高效的稀疏矩阵格式存储。Numba 库通过即时编译 (JIT) Python 代码为机器码，可以显著加速包含循环和条件判断的数值计算。SciPy 提供的 csr_matrix（Compressed Sparse Row matrix，压缩稀疏行矩阵）则非常适合存储行内非零元素分布不均的稀疏数据。

以下是结合 Numba 和 CSR 矩阵实现高效稀疏交叉距离计算的详细步骤和代码：

1. 欧几里得距离计算函数

首先，我们定义一个 Numba 加速的欧几里得距离计算函数。经验表明，在 Numba 环境下，手动实现的循环计算欧几里得距离通常比调用 np.linalg.norm 更快。

import numba as nb
import numpy as np
import scipy
import math

@nb.njit()
def euclidean_distance(vec_a, vec_b):
    """
    计算两个向量之间的欧几里得距离。
    使用 Numba 加速，避免 Python 循环开销。
    """
    acc = 0.0
    for i in range(vec_a.shape[0]):
        acc += (vec_a[i] - vec_b[i]) ** 2
    return math.sqrt(acc)

@nb.njit() 装饰器指示 Numba 在函数首次调用时将其编译为优化的机器码，从而实现接近 C 语言的执行速度。

2. Numba 加速的稀疏距离填充核心函数

接下来是算法的核心部分，一个 Numba 加速的函数，负责遍历掩码，条件性地计算距离，并填充稀疏矩阵所需的数据结构 (data, indicies, indptr)。

@nb.njit()
def masked_distance_inner(data, indicies, indptr, matrix_a, matrix_b, mask):
    """
    Numba 加速的核心函数，根据掩码条件性地计算距离，
    并填充 CSR 矩阵的 data, indicies, indptr 数组。
    """
    write_pos = 0 # 记录当前写入稀疏数据的位置
    N, M = matrix_a.shape[0], matrix_b.shape[0]

    # 遍历所有可能的 A 和 B 向量对
    for i in range(N):
        for j in range(M):
            if mask[i, j]: # 如果掩码指示需要计算此距离
                # 记录距离值
                data[write_pos] = euclidean_distance(matrix_a[i], matrix_b[j])
                # 记录此距离对应的列索引
                indicies[write_pos] = j
                write_pos += 1
        # 记录当前行结束时，data/indicies 中元素的数量
        # indptr[i+1] 指示第 i 行在 data/indicies 中结束的位置
        indptr[i + 1] = write_pos

    # 调试断言，确保所有预分配的内存都被使用
    assert write_pos == data.shape[0]
    assert write_pos == indicies.shape[0]
    # data, indicies, indptr 通过参数修改，无需返回

这个函数直接操作 NumPy 数组，避免了 Python 对象创建和管理带来的开销。indptr 数组的构建方式是 CSR 矩阵的关键：indptr[k] 存储的是第 k 行之前所有非零元素的总数，indptr[k+1] - indptr[k] 则表示第 k 行的非零元素数量。

3. 稀疏距离计算主函数

最后，我们封装一个 Python 函数，负责设置稀疏矩阵的结构（如预计算非零元素的总数），调用 Numba 核心函数，并最终构造并返回 scipy.sparse.csr_matrix 对象。

def masked_distance(matrix_a, matrix_b, mask):
    """
    计算并返回根据掩码过滤的稀疏成对距离矩阵。
    """
    N, M = matrix_a.shape[0], matrix_b.shape[0]
    assert mask.shape == (N, M), "掩码形状必须与 A 和 B 矩阵的维度兼容。"

    # 确保掩码是布尔类型
    mask = mask != 0

    # 计算稀疏矩阵中非零元素的总数，用于预分配内存
    sparse_length = mask.sum()

    # 为 CSR 矩阵的数据、列索引和行指针预分配内存
    # 注意：这些数组无需零初始化，因为它们将在 Numba 函数中被完全填充
    data = np.empty(sparse_length, dtype='float64')    # 存储距离值
    indicies = np.empty(sparse_length, dtype='int64') # 存储列索引
    indptr = np.zeros(N + 1, dtype='int64')             # 存储行指针

    # 调用 Numba 加速的核心函数进行计算和填充
    masked_distance_inner(data, indicies, indptr, matrix_a, matrix_b, mask)

    # 使用填充好的数据构建 CSR 稀疏矩阵
    return scipy.sparse.csr_matrix((data, indicies, indptr), shape=(N, M))

4. 完整示例与性能基准测试

为了验证其效率，我们创建一个更大的随机数据集进行测试：

# 示例数据
A_big = np.random.rand(2000, 10)
B_big = np.random.rand(4000, 10)
# 创建一个非常稀疏的掩码，非零元素比例小于 0.1%
M_big = np.random.rand(A_big.shape[0], B_big.shape[0]) < 0.001

# 使用 %timeit 进行性能测试 (在 IPython/Jupyter 环境中运行)
# %timeit masked_distance(A_big, B_big, M_big)
# 示例输出: 13.5 ms ± 66.6 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

# 对比原始方法（如果内存允许）
# diff_big = A_big[:,None] - B_big[None,:]
# %timeit np.linalg.norm(diff_big, ord=2, axis=2) * M_big
# 示例输出: 556 ms ± 3.74 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

通过上述基准测试，我们可以观察到显著的性能提升。在 A_big 形状为 (2000, 10)，B_big 形状为 (4000, 10)，且掩码 M_big 稀疏度为 0.1% 的情况下，Numba 加速的稀疏方法比原始的 NumPy 广播方法快了约 40 倍。对于更长的向量（例如，向量维度为 1000），性能提升甚至可以达到 1000 倍。这种巨大的性能差距主要来源于避免了大量不必要的计算和高效的内存管理。

注意事项与优化建议

1. Numba 的作用: Numba 的 njit 装饰器是实现高性能的关键。它将 Python 循环和数值操作编译为高效的机器码，消除了 Python 解释器的开销。
2. CSR 矩阵的优势: scipy.sparse.csr_matrix 是一种非常适合存储稀疏数据的格式，它只存储非零元素及其对应的索引，极大地节省了内存。对于大规模稀疏矩阵，这是必不可少的。
3. 数据类型优化:
   • 如果距离精度要求不高，可以将 float64 替换为 float32 (np.float32)，这可以减少内存占用并可能提高计算速度。
   • 如果矩阵的维度（行数、列数）和非零元素的总数小于 2³¹，可以将 int64 替换为 int32 (np.int32)，进一步节省内存。
4. 正确性验证: 在优化代码后，务必使用 np.allclose() 等方法与原始（但低效）的实现进行结果比对，确保新算法的正确性。
5. 稀疏度影响: 性能提升的幅度与掩码的稀疏度直接相关。掩码越稀疏（即需要计算的距离越少），性能提升越显著。如果掩码非常稠密，接近全连接，那么 Numba 稀疏方法可能不会带来显著优势，甚至可能因为额外的稀疏结构管理开销而略慢于纯 NumPy 广播。

总结

当需要计算两组向量间具有高度稀疏性的成对距离时，直接使用 NumPy 广播计算所有距离再进行掩码筛选的方法效率极低。通过结合 Numba 的即时编译能力和 SciPy 的 csr_matrix 稀疏数据结构，我们可以构建一个高度优化的解决方案。该方案通过有条件地计算所需距离并以内存高效的稀疏格式存储结果，显著提升了大规模数据集的处理速度和内存效率，是处理此类问题的专业且高效的选择。