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

Question

我正在尝试 3D 中的 plot 正态分布。 我有一个用 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))

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

有人可以帮我画出这个 3d plot 在 Python 吗？

Answer 1

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, 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()