Matplotlib 3D plot - 2D format for input data?

Question

I am plotting a function of two parameters with matplotlib. I copied an example in matplotlib tutorial and transformed with my own input data: vectors X and Y (equally spaces numbers in -3:3) and Z=peaks(X,Y) with peaks a function that I defined befohand. What is wrong?

def peaks(x,y):
   xsq=x**2
   ysq=y**2
   xsq_one=(x+1)**2
   ysq_one=(y+1)**2
   m1=3*(1-x)**2
   m2=10*(x/5-x**3-y**5)
   m3=1/3
   return m1*numpy.exp(-xsq-ysq_one)-m2*numpy.exp(-xsq-ysq)-m3*numpy.exp(-xsq_one-ysq)


from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.gca(projection='3d')
X=Y=numpy.arange(-3,3,0.01).tolist()
Z=[]
for i in range(len(X)):
Z.append(peaks(X[i],Y[i]))

ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
cset = ax.contour(X, Y, Z, zdir='z', offset=-100)
cset = ax.contour(X, Y, Z, zdir='x', offset=-40)
cset = ax.contour(X, Y, Z, zdir='y', offset=40)

ax.set_xlabel('X')
ax.set_xlim(-40, 40)
ax.set_ylabel('Y')
ax.set_ylim(-40, 40)
ax.set_zlabel('Z')
ax.set_zlim(-100, 100)

plt.show()

Thanks for advice!

Answer 1

You need to generate the meshgrid. X,Y and Z must be 2D arrays

import numpy
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d

def peaks(x,y):
    return x * numpy.sin(y)

fig = plt.figure()
ax = fig.gca(projection='3d')
X = Y= numpy.arange(-3, 3, 0.1).tolist()
X, Y = numpy.meshgrid(X, Y)

Z = []
for i in range(len(X)):
    Z.append(peaks(X[i],Y[i]))

# Z must be an array
Z = numpy.array(Z)

ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
cset = ax.contour(X, Y, Z, zdir='z', offset=-8)
cset = ax.contour(X, Y, Z, zdir='x', offset=-8)
cset = ax.contour(X, Y, Z, zdir='y', offset=8)

ax.set_xlabel('X')
ax.set_xlim(-8, 8)
ax.set_ylabel('Y')
ax.set_ylim(-8, 8)
ax.set_zlabel('Z')
ax.set_zlim(-8, 8)

plt.show()

Answer 2

The accepted answer no longer works. Sadly reviewers rejected my suggested edit which would have made it a working asnwer. So here is the same answer again but with the small change necessary to make it work in the current release of matplotlib.

    import numpy
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import axes3d

    def peaks(x,y):
        return x * numpy.sin(y)

    fig = plt.figure()
    ax = fig.gca(projection='3d')
    X = Y= numpy.arange(-3, 3, 0.1).tolist()
    X, Y = numpy.meshgrid(X, Y)

    Z = numpy.zeros(X.shape)
    for i in range(len(X)):
        for j in range(len(Y)):
            Z[i,j] = peaks(X[i,j],Y[i,j])

    ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
    cset = ax.contour(X, Y, Z, zdir='z', offset=-8)
    cset = ax.contour(X, Y, Z, zdir='x', offset=-8)
    cset = ax.contour(X, Y, Z, zdir='y', offset=8)

    ax.set_xlabel('X')
    ax.set_xlim(-8, 8)
    ax.set_ylabel('Y')
    ax.set_ylim(-8, 8)
    ax.set_zlabel('Z')
    ax.set_zlim(-8, 8)

    plt.show()