优化Python中稀疏向量对欧氏距离计算的性能

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)

print("原始方法计算的距离矩阵：\n", masked_distances)

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

@nb.njit()
def euclidean_distance(vec_a, vec_b):
    """
    计算两个向量之间的欧氏距离。
    使用Numba进行JIT编译以提高性能。
    """
    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矩阵的data、indices、indptr数组。
    """
    write_pos = 0  # 当前写入data和indices数组的位置
    N, M = matrix_a.shape[0], matrix_b.shape[0]

    for i in range(N):
        for j in range(M):
            if mask[i, j]:
                # 如果掩码为True，则计算距离并记录
                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]
    # data、indices、indptr数组通过参数修改，无需返回

def masked_distance(matrix_a, matrix_b, mask):
    """
    计算两组向量间由掩码指定的稀疏欧氏距离，并返回一个CSR稀疏矩阵。
    """
    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')    # 存储距离值
    indices = np.empty(sparse_length, dtype='int64')   # 存储列索引
    indptr = np.zeros(N + 1, dtype='int64')             # 存储每行的起始索引

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

    # 使用填充好的数组构建scipy.sparse.csr_matrix
    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

# 运行基准测试
# %timeit masked_distance(A_big, B_big, M_big)