NumPy 中高效向量化矩阵与向量幂序列的逐行加权求和

本文介绍如何使用 `np.einsum` 或广播机制,对二维数组按行应用不同权重向量(如幂级数),实现无显式循环的批量加权求和运算。

在科学计算中,常需将一个形状为 (m, n) 的矩阵 a 的每一行,分别与一个形状为 (m, k) 的权重矩阵 c 的对应行进行逐元素相乘,再沿行方向(即 axis=0)求和,最终得到形状为 (n, k) 的结果。例如,给定:

import numpy as np

a = np.array([[20, 12,  6],
              [12, 24, 18],
              [ 0, 14, 30]])
b = np.array([1, 0.5])
c = np.array([b ** i for i in range(0, 3)][::-1])  # shape: (3, 2)
# c = [[1.  , 0.25],
#      [1.  , 0.5 ],
#      [1.  , 1.  ]]

目标是计算:对 c 的每一列 j,执行 (a[i, :] * c[i, j]).sum(axis=0),即每行用 c 的第 i 行第 j 列加权,再按列累加。等价于矩阵乘法的“行-列加权点积”泛化形式。

推荐方案:np.einsum(清晰、高效、内存友好)
einsum 可精确描述索引操作:'ij,ik->jk' 表示对 a[i,j] 和 c[i,k] 沿共享索引 i 求和,输出 j(列维)× k(权重维):

result = np.einsum('ij,ik->jk', a, c)
print(result)
# [[32.   11. ]
#  [50.   29. ]
#  [54.  40.5]]

注意:np.einsum('ij,ik', a, c) 省略 ->jk 时默认求和并返回 jk,效果相同,但显式写出更利于理解。

⚠️ 替代方案:广播 + sum()(直观但内存开销大)
通过维度扩展实现广播:

# a[:, None] → (3, 1, 3)
# c[:, :, None] → (3, 2, 1)
# 广播后 shape: (3, 2, 3)
intermediate = a[:, None] * c[:, :, None]
result = intermediate.sum(axis=0)  # 沿 axis=0(原行维)求和 → (2, 3)
# [[32.  50.  54. ]
#  [11.  29.  40.5]]

⚠️ 此结果形状为 (k, n)(即 (2, 3)),与 einsum 的 (n, k) 转置关系。若需严格一致,可写为 .sum(0).T 或调整广播维度顺序。

? 关键总结

  • np.einsum('ij,ik->jk', a, c) 是最自然、高效且不易出错的解法,语义明确,性能优;
  • 广播方案虽直观,但会临时创建 (m, k, n) 大数组,当 m, k, n 较大时易触发内存瓶颈;
  • 避免使用 Python 循环或 np.vectorize(后者不真正向量化);
  • 若 c 来源于多项式基(如 b**0, b**1, b**2),还可考虑 np.polynomial.polynomial.polyvander 配合 @ 运算优化,但通用性不如 einsum。

掌握 einsum 的索引逻辑,是解锁 NumPy 高阶张量运算的关键一步。