如何用轴方向矩阵乘法替代循环实现W.T[i,:]@Sigma@W[:,i]计算？

你要计算的是W矩阵每一列对应的二次型W.T[i,:] @ Sigma @ W[:,i]，完全可以通过向量化操作替代循环，以下是几种高效实现方式：

方法1：矩阵乘法取对角线 #

原循环的计算等价于矩阵W.T @ Sigma @ W的对角线元素，直接取对角线即可：

result = np.diag(W.T @ Sigma @ W)

原理：W.T @ Sigma @ W生成一个m×m矩阵，其第(i,i)位置的元素正好对应原循环中第i次迭代的结果。不过该方法会额外计算矩阵的非对角线元素，当m较大时会有性能损耗。

方法2：逐元素相乘求和（更高效） #

先计算Sigma @ W，再与原W逐元素相乘后按行求和，避免多余计算：

result = np.sum(W * (Sigma @ W), axis=0)

原理：Sigma @ W的每一列是Sigma与W对应列的乘积，W和该矩阵逐元素相乘后，每列的累加和就是对应列的二次型结果。这种方法无需生成完整的m×m矩阵，性能比方法1更优。

方法3：Einstein求和约定（精准高效） #

用np.einsum直接描述求和规则，精准计算所需结果，无冗余运算：

result = np.einsum('ai,ab,bi->i', W, Sigma, W)

原理：ai对应W的第a行第i列元素，ab对应Sigma的第a行第b列元素，bi对应W的第b行第i列元素。通过收缩a和b维度，直接得到每个i对应的二次型结果，是三种方法中性能最优的。

验证示例 #

可以用以下代码验证三种方法与原循环结果一致：

import numpy as np

n, m = 3, 4
Sigma = np.random.rand(n, n)
W = np.random.rand(n, m)

# 原循环结果
loop_result = np.array([W.T[i,:] @ Sigma @ W[:,i] for i in range(m)])

# 方法1结果
diag_result = np.diag(W.T @ Sigma @ W)

# 方法2结果
sum_result = np.sum(W * (Sigma @ W), axis=0)

# 方法3结果
einsum_result = np.einsum('ai,ab,bi->i', W, Sigma, W)

# 验证所有结果一致
print(np.allclose(loop_result, diag_result))  # 输出True
print(np.allclose(loop_result, sum_result))  # 输出True
print(np.allclose(loop_result, einsum_result))  # 输出True

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