如何修改博客代码的einsum实现多协方差矩阵的多元正态logpdf向量化计算？

云原生阿岚

2026-6-12

向量化计算多样本多元正态分布的logpdf

假设你的变量形状如下：

cov: (n_samples, 2, 2)：每个样本对应的2维协方差矩阵
mu: (n_samples, 2)：每个样本对应的2维均值（若均值为固定值，可调整为对应形状适配广播）
pos: (m_points, 2)：已展平后的坐标点数组

以下是向量化实现的代码，重点调整了einsum部分适配你的数组形状：

import numpy as np

# 1. 计算协方差矩阵的行列式和逆矩阵
det_cov = np.linalg.det(cov)
inv_cov = np.linalg.inv(cov)

# 2. 构造每个点与对应均值的差向量（广播实现）
diff = pos[:, None, :] - mu[None, :, :]  # 最终形状: (m_points, n_samples, 2)

# 3. 计算二次型项：(x-mu)^T @ inv_cov @ (x-mu)
quadratic = np.einsum('mni,nij,mnj->mn', diff, inv_cov, diff)

# 4. 组合所有项得到logpdf
d = 2  # 数据维度
log_pdf = -0.5 * d * np.log(2 * np.pi) - 0.5 * np.log(det_cov)[None, :] - 0.5 * quadratic

关键部分说明：

广播构造差向量：通过pos[:, None, :]将点数组扩展为(m_points, 1, 2)，mu[None, :, :]扩展为(1, n_samples, 2)，广播后相减得到每个点对应每个样本的差向量。
einsum适配3D协方差：'mni,nij,mnj->mn'的含义是：
- mni：遍历每个点(m)、每个样本(n)、每个维度(i)
- nij：遍历每个样本(n)、输入维度(i)、输出维度(j)（对应协方差逆矩阵的维度）
- 先完成diff与inv_cov的矩阵乘法（得到mnj形状的中间结果），再与diff的mnj维度做内积求和，最终得到每个点-样本对的二次型值（mn形状）。