如何高效对多组坐标计算x^T A x形式的链式矩阵乘法？

问题描述 #

需要批量计算一组n维向量的二次型 xᵀAx，其中x是每行对应一个n维坐标的矩阵，A是n×n方阵。当前通过Python循环实现，希望找到更高效的向量化方法。

示例场景（二维向量）：

import numpy as np
x = np.column_stack([[1,2,3,4,5],[6,7,8,9,0]])
A = np.array([[1,0],[0,2]])

print(x[0] @ A @ x[0])  # 单个向量计算结果正确
# 目标：高效生成所有x[i] @ A @ x[i]的数组
y_loop = [x[i] @ A @ x[i] for i in range(x.shape[0])]  # 循环实现，效率较低

高效解决方案 #

方法1：向量化矩阵运算（最优性能） #

通过矩阵乘法结合逐元素相乘与轴求和，完全避免循环，利用numpy的底层优化提升效率：

# 先计算x与A的乘积
xA = x @ A
# 逐元素相乘后沿每行求和，得到每个向量的二次型结果
y = np.sum(xA * x, axis=1)

原理：xA[i]等价于x[i]@A，xA[i] * x[i]是逐元素对应相乘，对该行求和就等于x[i]@A@x[i]，所有操作均为向量化运算，效率远高于循环。

方法2：爱因斯坦求和约定（直观清晰） #

使用np.einsum可以直接通过索引描述二次型的计算逻辑，代码可读性强：

# 完整索引写法：i代表每个向量，j/k代表维度
y = np.einsum('ij,jk,ik->i', x, A, x)
# 简化版：先计算x@A，再与x做行内点积
y = np.einsum('ik,ik->i', x @ A, x)

方法3：矩阵点积取对角线（代码简洁） #

通过计算矩阵点积后的对角线元素得到结果，适合快速编写代码，但会额外计算非必要的元素，当向量数量较多时性能略逊于前两种方法：

y = np.dot(x @ A, x.T).diagonal()

原理：x@A是(m,n)矩阵，x.T是(n,m)矩阵，两者点积得到(m,m)矩阵，其对角线元素正好是每个x[i]@A@x[i]的结果。

验证结果 #

运行上述方法，结果与循环实现完全一致：

y_loop = np.array([x[i] @ A @ x[i] for i in range(x.shape[0])])
y_vec = np.sum((x@A)*x, axis=1)
print(np.array_equal(y_loop, y_vec))  # 输出 True

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