使用 Numba 和 CSR 矩阵高效计算稀疏交叉距离

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)

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)

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()
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 通过参数修改，无需返回

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))

# 示例数据
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)