Displaying multiple masks in different colours in pylab

An answer to this question on Stack Overflow.

Question

I've an array that includes decent observations, irrelevant observations (that I would like to mask out), and areas where there are no observations (that i would also like to mask out). I want to display this array as an image (using pylab.imshow) with two separate masks, where each mask is shown in a different colour.

I've found code for a single mask (here) in a certain colour, but nothing for two different masks:

masked_array = np.ma.array (a, mask=np.isnan(a))
cmap = matplotlib.cm.jet
cmap.set_bad('w',1.)
ax.imshow(masked_array, interpolation='nearest', cmap=cmap)

If possible, I'd like to avoid having to use a heavily distorted colour map but accept that that is an option.

Answer

In order to color some pixels red and others green, as appears in the the following image, I used the code below. (See code comments for details.)

Example of two masks being used with a matplotlib matrix graph

import numpy as np              #Used for holding and manipulating data
import numpy.random             #Used to generate random data
import matplotlib as mpl        #Used for controlling color
import matplotlib.colors        #Used for controlling color as well
import matplotlib.pyplot as plt #Use for plotting
#Generate random data
a = np.random.random(size=(10,10))
#This 30% of the data will be red
am1 = a<0.3                                 #Find data to colour special
am1 = np.ma.masked_where(am1 == False, am1) #Mask the data we are not colouring
#This 10% of the data will be green
am2 = np.logical_and(a>=0.3,a<0.4)          #Find data to colour special
am2 = np.ma.masked_where(am2 == False, am2) #Mask the data we are not colouring
#Colourmaps for each special colour to place. The left-hand colour (black) is
#not used because all black pixels are masked. The right-hand colour (red or
#green) is used because it represents the highest z-value of the mask matrices
cm1 = mpl.colors.ListedColormap(['black','red'])
cm2 = mpl.colors.ListedColormap(['black','green'])
fig = plt.figure()                          #Make a new figure
ax = fig.add_subplot(111)                   #Add subplot to that figure, get ax
#Plot the original data. We'll overlay the specially-coloured data
ax.imshow(a,   aspect='auto', cmap='Greys', vmin=0, vmax=1)
#Plot the first mask. Values we wanted to colour (`a<0.3`) are masked, so they
#do not show up. The values that do show up are coloured using the `cm1` colour
#map. Since the range is constrained to `vmin=0, vmax=1` and a value of
#`cm2==True` corresponds to a 1, the top value of `cm1` is applied to all such
#pixels, thereby colouring them red.
ax.imshow(am1, aspect='auto', cmap=cm1, vmin=0, vmax=1);
ax.imshow(am2, aspect='auto', cmap=cm2, vmin=0, vmax=1);
plt.show()

1: https://i.sstatic.net/dErl8.png