高效计算稀疏交叉差分：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)

# 计算所有可能的差分矩阵 (3, 5, 2)
diff = A[:, None] - B[None, :]
# 计算所有可能的欧氏距离 (3, 5)
distances = np.linalg.norm(diff, ord=2, axis=2)
# 应用掩码，将不需要的距离置为0
masked_distances = distances * M

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

@nb.njit()
def euclidean_distance(vec_a, vec_b):
    """
    使用 Numba 加速计算两个向量之间的欧氏距离。
    """
    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, indices, indptr, matrix_a, matrix_b, mask):
    """
    Numba 加速的核心函数，根据掩码计算距离并填充 CSR 矩阵的组成部分。
    """
    write_pos = 0
    N, M = matrix_a.shape[0], matrix_b.shape[0]

    # 遍历所有可能的向量对
    for i in range(N):
        for j in range(M):
            # 仅在掩码为 True 时计算距离
            if mask[i, j]:
                # 记录距离值
                data[write_pos] = euclidean_distance(matrix_a[i], matrix_b[j])
                # 记录该距离值对应的列索引
                indices[write_pos] = j
                write_pos += 1
        # 记录当前行结束后，下一行在 data 和 indices 中的起始位置
        indptr[i + 1] = write_pos

    # 确保所有预分配的空间都被使用
    assert write_pos == data.shape[0]
    assert write_pos == indices.shape[0]

def masked_distance(matrix_a, matrix_b, mask):
    """
    封装函数，用于设置并执行稀疏距离计算，返回 CSR 矩阵。
    """
    N, M = matrix_a.shape[0], matrix_b.shape[0]
    assert mask.shape == (N, M)

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

    # 统计稀疏矩阵中非零元素的总数
    sparse_length = mask.sum()

    # 预分配 CSR 矩阵的组成部分数组
    # data 和 indices 不需要初始化为零，直接填充更快
    data = np.empty(sparse_length, dtype='float64')    # 存储距离值
    indices = np.empty(sparse_length, dtype='int64')   # 存储列索引
    indptr = np.zeros(N + 1, dtype='int64')             # 存储行指针，需初始化首位为0

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

    # 构建并返回 SciPy CSR 稀疏矩阵
    return scipy.sparse.csr_matrix((data, indices, indptr), shape=(N, M))

# 示例数据生成
A_big = np.random.rand(2000, 10) # 2000个向量，每个10维
B_big = np.random.rand(4000, 10) # 4000个向量，每个10维
# 创建一个非常稀疏的掩码，只有0.1%的元素为True
M_big = np.random.rand(A_big.shape[0], B_big.shape[0]) < 0.001

# 性能测试 (在 Jupyter/IPython 环境中运行)
# %timeit masked_distance(A_big, B_big, M_big)