从 MATLAB 代码中使用 Numpy、Matploblib 3D 说明正态分布

问题描述

我正在尝试以 3D 形式绘制正态分布。我有一段用 MATLAB 编写的代码，但是我用 Python 编写它失败了。

完整的MATLAB代码为：

dsig = 0.25;
dx = 0.5;
mu = 0;
[X,SIGMA] = meshgrid(-10:dx:10,1:dsig:5);
Z = exp(-(X-mu).^2./(2*SIGMA.^2))./sqrt(2*pi*SIGMA.^2);
waterfall(X,SIGMA,Z)
xlabel('x')
ylabel('\sigma')
zlabel('f(x)')

到目前为止，我尝试用 Python 编写的代码是：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

dsig = 0.25
dx = 0.5
mu = 0

X = np.linspace(-10,dx,10)
SIGMA = np.linspace(1,dsig,5)

X,SIGMA = np.meshgrid(X,SIGMA)
Z = 1/(np.sqrt(2*np.pi*SIGMA*SIGMA))*np.exp(-(x-mu)**2/(2*SIGMA*SIGMA))

这段代码一直给我一个错误。

有人能帮我用 Python 绘制这个 3d 图吗？

解决方法

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator,FormatStrFormatter
import math
import scipy.stats as stats

mu = 0
variance = 1
sigma = math.sqrt(variance)

x = np.linspace(mu - 3*sigma,mu + 3*sigma,100)
y = np.linspace(mu - 3*sigma,100)
x,y = np.meshgrid(x,y)
r = np.sqrt(x**2 + y**2)
z = stats.norm.pdf(r,mu,sigma)

fig = plt.figure()
ax = fig.gca(projection='3d') # get current axis

surf = ax.plot_surface(x,y,z,cmap=cm.coolwarm,linewidth=0,antialiased=False)

ax.set_zlim(0,0.3)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

plt.show()