Why is a surface plot of integers not displaying properly using matplotlib plot

I was trying to replicate the answer found

There are two things wrong with this plot: (1) the surfaces are a constant color and (2) the color bar doesn't seem to be referencing the data.

Following the advice

I think this suggests the integer data in the first image needs to be normalized before displaying properly, but, if it's possible, I would really like to make this plot without normalization; I want integer values to display like the second image. Also, I'm not sure why the color bar isn't connected to the color scale of the images themselves and could use advice on how to fix that. Ideally, the color bar to be connected to all three surfaces, not just the im surface.

Thanks in advance!

CodePudding user response：

First, you have to normalize your data. Then, you pass the normalized data into the colormap to create the face colors:

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

data = np.random.randint(0,6,size=(49,512,512))

# create a Normalize object with the correct range
norm = colors.Normalize(vmin=data.min(), vmax=data.max())
# normalized_data contains values between 0 and 1
normalized_data = norm(data)
# extract the appropriate values
z = 20
z1 = 21
z2 = 22
i = normalized_data[z,:,:]
i1 = normalized_data[z1,:,:]
i2 = normalized_data[z2,:,:]

x = y = np.arange(0, 512, 1)


# here are the x,y and respective z values
X, Y = np.meshgrid(x, y)
Z = z*np.ones(X.shape)
Z1 = z1*np.ones(X.shape)
Z2 = z2*np.ones(X.shape)

# create the figure, add a 3d axis, set the viewing angle
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.view_init(10,60)

# here we create the surface plot, but pass V through a colormap
# to create a different color for each patch
im  = ax.plot_surface(X, Y, Z, facecolors=cm.viridis(i))
ax.plot_surface(X, Y, Z1, facecolors=cm.viridis(i1))
ax.plot_surface(X, Y, Z2, facecolors=cm.viridis(i2))

# create a scalar mappable to create an appropriate colorbar
sm = cm.ScalarMappable(cmap=cm.viridis, norm=norm)
fig.colorbar(sm)