Python曲面图值错误

Question

我已经习惯了python，现在想制作一个3D表面图。 我有三个变量x，y，z和一个强度函数I = I（x，y，z）。我想挑选一个特定的切片z = 250并绘制x和y方向的dsitribution。 问题是，我不知道如何为z值选择切片。我试过的所有版本都结束了一些错误

这是我的

import pandas as pd
import numpy as np #NumPy
import scipy as sp #SciPy
import matplotlib as mpl #Matplotlib(2D/3D)
import matplotlib.pyplot as plt #Matplotlib's pyplot
from pylab import * #Matplotlib's pylab
from mpl_toolkits.mplot3d import Axes3D
%matplotlib inline


x = linspace(-4,4,50)     # units mm
y = linspace(-4,4,50)     # units mm


# define beam parameters
I_0 = 1e0                # intensity
lambda_1 = 800e-9        # wavelength
w_0 = 1.5                # beam waist
z_r = pi*w_0**2*1e-6/lambda_1 # Rayleigh length in units mm

z = linspace(-4,4,500)*z_r
# calculate intensity profile

X,Y,Z = meshgrid(x,y,z)

w = w_0 *sqrt(1+(Z/z_r)**2)

I = I_0*((w_0/w)**2)*exp(-2*(X**2+Y**2)/(w**2))

fig = plt.figure()
ax = fig.gca(projection='3d')

surf = ax.plot_surface(X, Y, I[:,:,250], cmap=cm.coolwarm,linewidth=0,antialiased=False)

show()

我收到错误：ValueError：形状不匹配：无法将对象广播为单个形状

我该如何正确地做到这一点？ 编辑：我使用python 3.6和jupyter notebook

Answer 1

numpy.meshgrid(x,y,z)创建三个3D数组。 Matplotlib将需要2D数组作为输入。因此，您需要另一组x和y数组进行绘图。这些也可以通过numpy.meshgrid(x,y)创建，或者您可以只切片现有数组。

surf = ax.plot_surface(X[:,:,0], Y[:,:,0], I[:,:,250], ...)

完整示例：

import numpy as np #NumPy
import matplotlib.pyplot as plt #Matplotlib's pyplot
from mpl_toolkits.mplot3d import Axes3D
#%matplotlib inline

x = np.linspace(-4,4,50)     # units mm
y = np.linspace(-4,4,50)     # units mm

# define beam parameters
I_0 = 1e0                # intensity
lambda_1 = 800e-9        # wavelength
w_0 = 1.5                # beam waist
z_r = np.pi*w_0**2*1e-6/lambda_1 # Rayleigh length in units mm

z = np.linspace(-4,4,500)*z_r


X,Y,Z = np.meshgrid(x,y,z)

w = w_0 *np.sqrt(1+(Z/z_r)**2)
# calculate intensity profile
I = I_0*((w_0/w)**2)*np.exp(-2*(X**2+Y**2)/(w**2))

fig = plt.figure()
ax = fig.gca(projection='3d')

surf = ax.plot_surface(X[:,:,0], Y[:,:,0], I[:,:,250], 
                       cmap=plt.cm.coolwarm,linewidth=0,antialiased=False)

plt.show()