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

Question

I was trying to replicate the answer found here with my own data, which happens to be a 3D numpy array of integers. I got close with the following code:

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

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

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

z = 20
i = data[z,:,:]

z1 = 21
i1 = data[z1,:,:]

z2 = 22
i2 = data[z2,:,:]

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

But this produces the plot below.

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 here, I found that (1) can be solved by replacing data with a set of random numbers data = np.random.random(size=(49,512,512)), which produces the below image.

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!

Answer 1

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)