How to plot 3D surface with X, Y, Z when Z is a list of list in Python?

Question

In my case, X is a range(0, 100), Y is a range(0, 10), Z is a list of list. Z has the same length as X, which is 100, and each element list inside of Z has the same dimension of Y.

Z = [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], ..., [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]].

I have the following code, but it does not work, it complains two or more arrays have incompatible dimensions on axis 1.

fig = plt.figure(figsize=(200, 6))
ax = fig.add_subplot(1, 2, 1, projection='3d')
ax.set_xticklabels(x_ax)
ax.set_yticklabels(y_ax)
ax.set_title("my title of chart")

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=False)

ax.set_zlim(0, 100)

fig.colorbar(surf, shrink = 0.5, aspect = 5)
plt.show()

I guess the error is due to the data structure of Z, how do I make a compatible structure with X, and Y? Thanks

Answer 1

Here is a basic 3D surface plotting procedure. It seems that your X and Y are just 1D arrays. However, X , Y , and Z have to be 2D arrays of the same shape. numpy.meshgrid function is useful for creating 2D mesh from two 1D arrays.

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

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

x = np.array(np.linspace(-2,2,100))
y = np.array(np.linspace(-2,2,10))
X,Y = np.meshgrid(x,y)

Z = X * np.exp(-X**2 - Y**2);

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=False)
fig.colorbar(surf, shrink = 0.5, aspect = 5)

plt.show()