如何将二维高斯分布的等高线图限制为仅显示3σ变异范围

核心原理 #

二维高斯分布的3σ置信区间（覆盖99.73%的样本概率）对应马氏距离平方满足如下条件：
$$d^2 = (x-\mu)^T \Sigma^{-1} (x-\mu) < \chi^2_{2, 0.9973} \approx 11.83$$
其中$\chi^2_{2, 0.9973}$是自由度为2的卡方分布99.73%分位数值，所有超出该阈值的区域直接屏蔽即可实现范围限制。

前置修正：原代码已知bug #

• 协方差矩阵计算错误：原代码中covMat3的计算逻辑有误，无法得到正确的协方差矩阵，可直接用numpy内置函数简化计算：

  covMat3 = np.cov(localPoints.T, bias=False)
  

  如果要手动计算，正确的无偏估计公式为：

  covMat3 = np.einsum('ijk, ikj->jk', lPoints - meanXY2, lPoints - meanXY2) / (count - 1)
  

• 概率密度公式错误：原代码的probab函数前置系数写反，二维高斯概率密度的前置系数为$1/(2\pi \sqrt{det(\Sigma)})$，修正后的函数为：

  def probab(mean, covMat, lPoints):
      lPoints = lPoints[..., np.newaxis] if lPoints.ndim == 2 else lPoints
      det = np.linalg.det(covMat)
      if det > 0:
          factor1 = 1 / (2 * np.pi * np.sqrt(det))
          factor2 = np.exp((-1/2)*np.einsum('ijk,jl,ilk->ik', lPoints - mean, np.linalg.inv(covMat), lPoints - mean))
          return factor1 * factor2
  

3σ裁剪具体实现 #

在计算得到概率值z2后，新增马氏距离判断逻辑，将超出3σ范围的概率值设为NaN，matplotlib绘制等高线时会自动忽略空值区域，修改后的核心代码段如下：

if np.linalg.det(covMat3) > 0:   # 仅协方差矩阵可逆时计算
    Xx = np.linspace(x[i], x[i + 1], 1000)
    Yy = np.linspace(y[j], y[j + 1], 1000)
    Xx,Yy = np.meshgrid(Xx, Yy)
    lPoints = np.vstack((Xx.flatten(), Yy.flatten())).T
    # 计算马氏距离平方
    inv_cov = np.linalg.inv(covMat3)
    diff = lPoints - meanXY2.flatten()
    mahalanobis_sq = np.einsum('ij,jk,ik->i', diff, inv_cov, diff)
    # 3σ阈值，也可以用scipy.stats.chi2.ppf(0.9973, df=2)计算得到精确值
    chi2_3sigma = 11.829
    # 计算概率
    z2 = probab(meanXY2, covMat3, lPoints)
    # 超出3σ范围的位置设为NaN
    z2[mahalanobis_sq > chi2_3sigma] = np.nan
    z2 = z2.reshape(1000, 1000)
    # 绘制等高线
    cs = plt.contourf(Xx, Yy, z2)
    plt.clabel(cs)

可选优化方案 #

如果不需要保留网格点的完整概率值，也可以直接限制等高线的绘制层级：先计算3σ阈值对应的概率值，再传入levels参数给contourf即可，不需要修改z2的原始值。

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