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

import numba as nb
import numpy as np
import scipy
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, indicies, indptr, matrix_a, matrix_b, mask):
    """
    Numba JIT编译的核心函数，用于根据掩码计算并填充稀疏矩阵的数据。

    参数:
        data (np.ndarray): 存储非零距离值的数组。
        indicies (np.ndarray): 存储非零距离值对应列索引的数组。
        indptr (np.ndarray): 存储每行在data和indicies中起始位置的数组。
        matrix_a (np.ndarray): 第一个向量集合。
        matrix_b (np.ndarray): 第二个向量集合。
        mask (np.ndarray): 布尔型掩码矩阵，指示哪些距离需要计算。
    """
    write_pos = 0  # 当前写入data和indicies的位置
    N, M = matrix_a.shape[0], matrix_b.shape[0]

    for i in range(N):  # 遍历 matrix_a 的每一行
        for j in range(M):  # 遍历 matrix_b 的每一行
            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] = write_pos

    # 断言所有预分配的内存都被使用
    assert write_pos == data.shape[0]
    assert write_pos == indicies.shape[0]

def masked_distance(matrix_a, matrix_b, mask):
    """
    计算并返回一个稀疏矩阵，其中包含根据掩码筛选出的成对欧氏距离。

    参数:
        matrix_a (np.ndarray): 第一个向量集合。
        matrix_b (np.ndarray): 第二个向量集合。
        mask (np.ndarray): 布尔型掩码矩阵，指示哪些距离需要计算。

    返回:
        scipy.sparse.csr_matrix: 包含所需距离的稀疏矩阵。
    """
    N, M = matrix_a.shape[0], matrix_b.shape[0]
    assert mask.shape == (N, M)

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

    # 计算稀疏矩阵将包含的非零元素总数
    sparse_length = mask.sum()

    # 预分配存储稀疏矩阵数据的数组
    # 注意：这些数组不需要初始化为零，Numba函数会直接写入
    data = np.empty(sparse_length, dtype='float64')    # 存储距离值
    indicies = np.empty(sparse_length, dtype='int64')  # 存储列索引
    indptr = np.zeros(N + 1, dtype='int64')            # 存储行指针，第一个元素为0

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

    # 构建并返回SciPy的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%的元素为True
M_big = np.random.rand(A_big.shape[0], B_big.shape[0]) < 0.001

# 使用 %timeit 魔法命令测量执行时间
# %timeit masked_distance(A_big, B_big, M_big)