如何用Numpy高效实现大规模向量的逐对正交投影

问题场景 #

现有两个规模相同的2D Numpy笛卡尔向量数组A、B，结构如下：

A = [[ax1, ay1, az1], [a2], [a3], ...]  # 每个元素是三维向量a_i
B = [[bx1, by1, bz1], [b2], [b3], ...]  # 每个元素是三维向量b_i

需要将A中每个向量正交投影到B中对应的向量上，得到输出数组C。单个向量的投影公式为：

c_i = (np.dot(a_i, b_i) / np.dot(b_i, b_i)) * b_i

但A、B各包含约5000万个向量，Python循环效率极低，直接替换2D数组或添加[:,None]修饰符未实现逐向量投影，需纯Numpy无循环的高效方案。

无循环实现方案 #

方案1：使用np.einsum计算逐向量点积 #

np.einsum能精准控制张量运算，高效完成逐向量点积计算：

import numpy as np

# 计算每对对应向量的点积A·B，结果为形状(N,)的一维数组
dot_ab = np.einsum('ij,ij->i', A, B)
# 计算每个B向量的模长平方B·B，结果为形状(N,)的一维数组
dot_bb = np.einsum('ij,ij->i', B, B)
# 将系数转为(N,1)形状，以便和B数组广播相乘
coeffs = dot_ab / dot_bb[:, np.newaxis]
# 逐向量相乘得到投影结果C，形状与A、B一致
C = coeffs * B

方案2：使用np.sum结合轴参数 #

若对einsum不熟悉，可通过np.sum指定求和轴实现相同效果：

# 逐向量计算A·B：对应元素相乘后沿轴1求和
dot_ab = np.sum(A * B, axis=1)
# 逐向量计算B·B：B向量元素平方后沿轴1求和
dot_bb = np.sum(B * B, axis=1)
# 调整系数维度实现广播，再与B相乘得到投影
C = (dot_ab / dot_bb)[:, np.newaxis] * B

核心原理 #

• 两种方案均先计算逐向量的点积与模长平方，得到一维数组（每个元素对应一对A、B向量的运算结果）
• 将模长平方数组转为(N,1)的二维形状，Numpy会自动将其广播为与B数组匹配的(N,3)形状，实现逐元素的系数乘法
• 完全依托Numpy的底层C优化实现，避免Python循环的性能损耗，可高效处理千万级规模的向量数据

内容的提问来源于stack exchange，提问作者Mark Huisjes