3D plot with an 2D array python matplotlib

Question

I have 2 1D arrays with the values of x and y, and also a 2D array with the values of z for each point where the columns correspond to the x values and the rows to the y values. Is there any way to get a plot_surface with this data? when I try to do it it returns me no plot. Here is the code: (calculate_R is a function I made for the program)

x=np.arange(0,10,1)
y=np.arange(0,1,0.2)
lx= len(x)
ly=len(y)

z=np.zeros((lx,ly))

for i in range(lx):
    for j in range(ly):
        z[i,j]=calculate_R(y[j],x[i])

fig = plt.figure()
ax = Axes3D(fig)
x, y = np.meshgrid(x, y)
ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap='hot')

Solution

You forgot to call plt.show() to display your plot.

Note that you might be able to exploit numpy vectorization to speed up the calculation of z:

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

x = np.arange(0,10,1)
y = np.arange(0,1,0.2)

xs, ys = np.meshgrid(x, y)
# z = calculate_R(xs, ys)
zs = xs**2 + ys**2

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(xs, ys, zs, rstride=1, cstride=1, cmap='hot')
plt.show()

Here, I used a simple function, since you didn't supply a fully working example.

Comments

[Karlo]

Is it possible to replace zs by an np.array (matrix mat) and x, y by np.arange(0, xsize, 1), np.arange(0, ysize, 1) where xsize, ysize = mat.shape, in order to obtain the 3D plot of a matrix?

[David Zwicker]

Yes, that should be possible.

[Karlo]

Indeed: see this question