3D discrete heatmap in matplotlib
<p>I have a list of tuples in python containing 3-dimenstional data, where each tuple is in the form: (x, y, z, data_value), i.e., I have data values at each (x, y, z) coordinate. I would like to make a 3D discrete heatmap plot where the colors represent the value of data_values in my list of tuples. Here, I give an example of such a heatmap for a 2D dataset where I have a list of (x, y, data_value) tuples:</p>

<pre><code>import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np
from random import randint

# x and y coordinates
x = np.array(range(10))
y = np.array(range(10,15))
data = np.zeros((len(y),len(x)))

# Generate some discrete data (1, 2 or 3) for each (x, y) pair
for i,yy in enumerate(y):
    for j, xx in enumerate(x):
        data[i,j] = randint(1,3)

# Map 1, 2 and 3 to 'Red', 'Green' qnd 'Blue', respectively
colormap = colors.ListedColormap(['Red', 'Green', 'Blue'])
colorbar_ticklabels = ['1', '2', '3']

# Use matshow to create a heatmap
fig, ax = plt.subplots()
ms = ax.matshow(data, cmap = colormap, vmin=data.min() - 0.5, vmax=data.max() + 0.5, origin = 'lower')

# x and y axis ticks
ax.set_xticklabels([str(xx) for xx in x])
ax.set_yticklabels([str(yy) for yy in y])
ax.xaxis.tick_bottom()

# Put the x- qnd y-axis ticks at the middle of each cell 
ax.set_xticks(np.arange(data.shape[1]), minor = False)
ax.set_yticks(np.arange(data.shape[0]), minor = False)

# Set custom ticks and ticklabels for color bar
cbar = fig.colorbar(ms,ticks = np.arange(np.min(data),np.max(data)+1))
cbar.ax.set_yticklabels(colorbar_ticklabels)

plt.show()
</code></pre>

<p>This generates a plot like this:
<a href="https://i.stack.imgur.com/rwH59.jpg" rel="noreferrer"><img src="https://i.stack.imgur.com/rwH59.jpg" alt="enter image description here"></a></p>

<p>How can I make a similar plot in 3D-space (i.e., having a z-axis), if my data have a third dimension. For example, if</p>

<pre><code># x and y and z coordinates
x = np.array(range(10))
y = np.array(range(10,15))
z = np.array(range(15,20))
data = np.zeros((len(y),len(x), len(y)))

# Generate some random discrete data (1, 2 or 3) for each (x, y, z) triplet. 
# Am I defining i, j and k correctly here?
for i,yy in enumerate(y):
    for j, xx in enumerate(x):
        for k, zz in enumerate(z):
            data[i,j, k] = randint(1,3)
</code></pre>

<p>I sounds like <a href="http://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html" rel="noreferrer">plot_surface in mplot3d</a> should be able to do this, but z in the input of this function is essentially the value of data at (x, y) coordinate, i.e., (x, y, z = data_value), which is different from what I have, i.e., (x, y, z, data_value). </p>