Python surface plot Value error

Question

I'm getting used to python and now wanna make a 3D surface plot. I have three variables x,y,z and an Intensity function I=I(x,y,z). I want to pick out a particular slice z=250 and plot the dsitribution for the x and y direction. The Problem is, I dont know how to pick a slice for a z-value. All Versions I tried ended iin some error

Here is what i have

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

I get an error: ValueError: shape mismatch: objects cannot be broadcast to a single shape

How do I do it correctly? Edit: I'm using python 3.6 with jupyter notebook

Answer 1

numpy.meshgrid(x,y,z) creates three 3D arrays. Matplotlib will need 2D arrays as input. You would hence need another set of x and y arrays for plotting. Those can also be created via numpy.meshgrid(x,y) or you can just slice the existing arrays.

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

Complete example:

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