Discretised color map with transparency in Matplotlib

Question

I'm looking to plot a matplotlib colormesh on top of another colormesh. The bottom plot is simply grayscale.

The one which sits on top should however draw a transparent square when the value of the passed array is 0, and a different color for each other number in the passed array. These are 2d numpy arrays.

Currently I have:

plt.pcolormesh(array1, vmin = -32, vmax = 32, cmap = plt.cm.binary)
plt.pcolormesh(array2, cmap = plt.cm.spectral)

Obviously this doesn't produce what I'm looking for, and I assume the way to do this is to generate my own colormap, I've read this guide: http://wiki.scipy.org/Cookbook/Matplotlib/ColormapTransformations but this doesn't seem to address transparency, nor how to make specific values map to specific colors.

As a short example of what I'd like, an array:

[[0, 1]
 [2, 3]]

Should produce a grid looking like:

[[transparent, red
 [green, yellow]]

How do I go about doing this? Merging the arrays together isn't an option, as the bottom dataset is a height map, and the values of this will likely always span the values of the second array (these are agent IDs).

Answer 1

This code should do something akin to what you require:

Edit using masked_array:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as colors
import numpy.ma as ma

#See http://stackoverflow.com/questions/18926031/how-to-extract-a-subset-of-a-colormap-as-a-new-colormap-in-matplotlib
def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
    new_cmap = colors.LinearSegmentedColormap.from_list(
        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
        cmap(np.linspace(minval, maxval, n)))
    return new_cmap

#truncate the colourmap
n_colours = 4
new_cmap = truncate_colormap(cm.get_cmap('spectral_r'), 0, 0.4, n=n_colours)

#discretise the colourmap
bounds = np.linspace(0,n_colors,n_colours+1)
norm = colors.BoundaryNorm(bounds, new_cmap.N)

#build array one
array1 = np.random.rand(10,10)

#build array two
array2 = np.random.randint(0,5,100).reshape(10,10)

#mask the array
array2 = ma.masked_array(array2, array2==0)

#plot it
plt.pcolormesh(array1,cmap = plt.cm.binary)
plt.pcolormesh(array2,cmap = new_cmap, norm=norm)
cbar = plt.colorbar()
plt.show()

Here is the new output using a masked array: