python - Change bar color in a 3D bar plot in matplotlib based on value

Question

I have a 3D bar plot in matplotlib which consists of a total 165 bars and at the moment it is quite chaotic.

.

I would like to change the colour of the bars based on the discreet z-values: 0,1,2.

I know there is the option to change colour bar in 1D bar plots based on specific values by using masks as in Color matplotlib bar chart based on value.

And there is also a question on how to change bar colour based on values: Defining colors of Matplotlib 3D bar plot

I am not sure If i perfectly comprehend the given answer but I cannot make it work in this case.

Code is:

   data = [[0 0 0 2 0 0 1 2 0 0 0]
            [0 0 2 2 0 0 0 0 2 0 0]
            [1 0 2 2 1 2 0 0 2 0 2]
            [1 0 2 2 0 2 0 2 2 2 2]
            [2 2 2 2 2 2 2 2 2 2 2]
            [2 2 0 2 2 2 2 2 2 2 2]
            [0 2 2 0 2 2 2 2 2 2 2]
            [1 2 0 0 2 1 2 2 0 0 2]
            [0 0 2 1 0 0 2 0 0 0 0]
            [2 1 2 2 0 0 0 2 0 0 2]
            [2 2 2 0 2 0 0 0 2 2 2]
            [2 2 0 0 2 2 2 2 2 0 0]
            [2 2 1 2 0 0 0 2 2 2 0]
            [2 0 0 2 0 0 2 2 2 2 2]
            [2 0 0 2 0 2 2 2 2 2 2]]

   ly = len(data[0])
   lx = len(data[:,0])
   xpos = np.arange(0,lx,1)    # Set up a mesh of positions
   ypos = np.arange(0,ly,1)
   xpos, ypos = np.meshgrid(xpos+0.25, ypos+0.25)

   xpos = xpos.flatten()   # Convert positions to 1D array
   ypos = ypos.flatten()
   zpos = np.zeros(lx*ly)

   dx = 0.5 * np.ones_like(zpos)
   dy = dx.copy()
   dz = data.flatten()


   ys = np.array([float(yi) for yi in y[1:]])

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

   # all blue bars
   #ax.bar3d(xpos,ypos,zpos, dx, dy, dz, color='b')

   # try changing color bars

   colors = ['r','g','b']
   for i in range(0,3):

       ax.bar3d(xpos[i], ypos[i], zpos[i], dx, dy, dz[i], alpha=0.1, 
                    color=colors[i])

   ax.set_xlabel('X')
   ax.set_ylabel('Y')
   ax.set_zlabel('Z')


plt.show()

Answer 1

As seen from the documentation of bar3d, color can be an array, with one color per bar.

This makes it quite easy to colorize all bars in a single call to bar3d; we just need to convert the data array to an array of colors which can be done using a colormap,

colors = plt.cm.jet(data.flatten()/float(data.max()))

(Note, that a colormap takes values between 0 and 1, so we need to normalize the values into this range.)

Complete example:

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

data = np.array([ [0, 0, 0, 2, 0, 0, 1, 2, 0, 0, 0],
         [0, 0, 2, 2, 0, 0, 0, 0, 2, 0, 0],
         [1, 0, 2, 2, 1, 2, 0, 0, 2, 0, 2],
         [1, 0, 2, 2, 0, 2, 0, 2, 2, 2, 2],
         [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
         [2, 2, 0, 2, 2, 2, 2, 2, 2, 2, 2],
         [0, 2, 2, 0, 2, 2, 2, 2, 2, 2, 2],
         [1, 2, 0, 0, 2, 1, 2, 2, 0, 0, 2],
         [0, 0, 2, 1, 0, 0, 2, 0, 0, 0, 0],
         [2, 1, 2, 2, 0, 0, 0, 2, 0, 0, 2],
         [2, 2, 2, 0, 2, 0, 0, 0, 2, 2, 2],
         [2, 2, 0, 0, 2, 2, 2, 2, 2, 0, 0],
         [2, 2, 1, 2, 0, 0, 0, 2, 2, 2, 0],
         [2, 0, 0, 2, 0, 0, 2, 2, 2, 2, 2],
         [2, 0, 0, 2, 0, 2, 2, 2, 2, 2, 2]])


ypos, xpos  = np.indices(data.shape) 

xpos = xpos.flatten()   
ypos = ypos.flatten()
zpos = np.zeros(xpos.shape)

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

colors = plt.cm.jet(data.flatten()/float(data.max()))
ax.bar3d(xpos,ypos,zpos, .5,.5,data.flatten(), color=colors)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()