绘制二维高斯分布的python代码

要在Python中绘制二维高斯分布（正态分布），你可以使用matplotlib和numpy库来生成和可视化数据。以下是一个完整的示例，说明如何绘制均值为 $\mu =(0,0)$ 并且具有某个协方差矩阵 $\Sigma$ 的二维高斯分布的等高线图。

示例代码

import numpy as np
import matplotlib.pyplot as plt

# 设置均值和协方差矩阵
mu = np.array([0, 0])
Sigma = np.array([[1, 0.5], [0.5, 1]])  # 例如，[1, 0.5; 0.5, 1] 协方差矩阵

# 创建网格点
x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)

# 将网格点组合成二维正态分布的输入形式
Z = np.dstack((X, Y))

# 计算二维高斯分布的概率密度函数
def multivariate_gaussian(pos, mu, Sigma):
    n = mu.shape[0]
    Sigma_det = np.linalg.det(Sigma)
    Sigma_inv = np.linalg.inv(Sigma)
    N = np.sqrt((2*np.pi)**n * Sigma_det)
    fac = np.einsum('...k,kl,...l->...', pos-mu, Sigma_inv, pos-mu)
    return np.exp(-fac / 2) / N

# 计算每个网格点的概率密度
pdf = multivariate_gaussian(Z, mu, Sigma)

# 绘制等高线图
plt.contour(X, Y, pdf, colors='black', alpha=0.5)
plt.contourf(X, Y, pdf, cmap='Blues')
plt.colorbar()
plt.title('2D Gaussian distribution')
plt.xlabel('X')
plt.ylabel('Y')
plt.axis('equal')
plt.show()

代码解释

1. 设置均值和协方差：定义二维高斯分布的均值和协方差矩阵。
2. 创建网格点：通过numpy.linspace和numpy.meshgrid生成一个均匀的网格，用于计算和显示二维高斯分布的概率密度。
3. 定义概率密度函数：multivariate_gaussian函数通过协方差矩阵和均值来计算位置数组pos上的概率密度。
4. 绘图：使用matplotlib.pyplot的contour和contourf函数绘制概率密度的等高线和填充等高线图。colorbar添加了一个颜色条，用于解释颜色与概率密度的关系。

这个代码段将生成一个展示所选均值和协方差参数下二维高斯分布的图形，帮助直观理解该分布在二维空间中的形状。