扩展维度后如何用np.einsum向量化计算平方马氏距离？

代码实现 #

首先构造维度匹配的差值矩阵，再通过np.einsum完成向量化计算：

import numpy as np

# 假设X(m,n), U(l,k,n), T(l,k,n,n)已定义
diff = X[np.newaxis, np.newaxis, :, :] - U[:, :, np.newaxis, :]
D = np.einsum('abcd,abde,abce->acb', diff, T, diff)

维度解析 #

1. 差值矩阵构造：
   • 将X扩展为(1,1,m,n)，U扩展为(l,k,1,n)，相减后得到diff的形状为(l,k,m,n)，对应字母标记：a=l（组维度）、b=k（每个组内的中心维度）、c=m（样本维度）、d=n（特征维度）。
2. einsum维度映射：
   • diff的维度：a,b,c,d（对应l,k,m,n）
   • T的维度：a,b,d,e（对应l,k,n,n，协方差矩阵的两个特征维度）
   • 第二个diff的维度：a,b,c,e（和第一个diff一致，特征维度用e标记以区分）
   • 输出acb对应维度l,m,k，正好匹配你需要的结果形状。

验证逻辑 #

这个计算等价于对每个组a∈[0,l)、每个中心b∈[0,k)、每个样本c∈[0,m)，计算：
$$D[a,c,b] = \text{diff}[a,b,c,:]^T \cdot T[a,b,:,:] \cdot \text{diff}[a,b,c,:]$$
完全符合平方马氏距离的定义，且全程向量化，效率和原代码保持一致。

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