高效计算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)
# 计算所有成对距离（L2范数）
distances = np.linalg.norm(diff, ord=2, axis=2)                     # (3, 5)
# 应用掩码，保留所需距离
masked_distances = distances * M                                    # (3, 5)

print("计算的距离矩阵：\n", distances)
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 加速，避免 np.linalg.norm 的开销。
    """
    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 (np.ndarray): 存储非零距离值的数组。
        indicies (np.ndarray): 存储非零距离值对应列索引的数组。
        indptr (np.ndarray): 存储每行在 data/indicies 中起始位置的数组。
        matrix_a (np.ndarray): 第一个向量集。
        matrix_b (np.ndarray): 第二个向量集。
        mask (np.ndarray): 布尔掩码矩阵，指示哪些距离需要计算。
    """
    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])
                # 记录该距离值对应的列索引
                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]
    # data, indicies, indptr 会在函数外部被修改并用于构建 CSR 矩阵

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

    # 为 CSR 矩阵预分配内存。这些数组不需要初始化为零，直接分配内存更高效。
    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% 的元素为 True）
M_big = np.random.rand(A_big.shape[0], B_big.shape[0]) < 0.001

# 使用优化的方法计算稀疏距离
sparse_distances = masked_distance(A_big, B_big, M_big)

print(f"稀疏距离矩阵的形状: {sparse_distances.shape}")
print(f"稀疏距离矩阵的非零元素数量: {sparse_distances.nnz}")
print(f"稀疏距离矩阵的密度: {sparse_distances.nnz / (sparse_distances.shape[0] * sparse_distances.shape[1]):.6f}")

# 性能基准测试 (在Jupyter/IPython环境中运行)
# %timeit masked_distance(A_big, B_big, M_big)
# 
# 原始方法的性能基准测试 (仅供参考，不推荐在生产环境运行大型矩阵)
# %timeit np.linalg.norm(A_big[:,None] - B_big[None,:], ord=2, axis=2) * M_big