Plotting masked numpy array leads to incorrect colorbar

Question

I'm trying to create a custom color bar for a matplotlib PolyCollection. Everything seems ok until I attempt to plot a masked array. The color bar no longer shows the correct colors even though the plot does. Is there a different procedure for plotting masked arrays?

I'm using matplotlib 1.4.0 and numpy 1.8.

Here's my plotting code:

import numpy
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection

vertices = numpy.load('vertices.npy')
array = numpy.load('array.npy')

# Take 2d slice out of 3D array
slice_ = array[:, :, 0:1].flatten(order='F')

fig, ax = plt.subplots()

poly = PolyCollection(vertices, array=slice_, edgecolors='black', linewidth=.25)

cm = mpl.colors.ListedColormap([(1.0, 0.0, 0.0), (.2, .5, .2)])
poly.set_cmap(cm)

bounds = [.1, .4, .6]
norm = mpl.colors.BoundaryNorm(bounds, cm.N)

fig.colorbar(poly, ax=ax, orientation='vertical', boundaries=bounds, norm=norm)
ax.add_collection(poly, autolim=True)
ax.autoscale_view()
plt.show()

Here's what the plot looks like:

However, when I plot a masked array with the following change before the slicing:

array = numpy.ma.array(array, mask=array > .5)

I get a color bar that now shows only a single color. Even though both colors are (correctly) still shown in the plot.

Is there some trick to keeping a colobar consistent when plotting a masked array? I know I can use cm.set_bad to change the color of masked values, but that's not quite what I'm looking for. I want the color bar to show up the same between these two plots since both colors and the color bar itself should remain unchanged.

Answer 1

Pass the BoundaryNorm to the PolyCollection, poly. Otherwise, poly.norm gets set to a matplotlib.colors.Normalize instance by default:

In [119]: poly.norm
Out[119]: <matplotlib.colors.Normalize at 0x7faac4dc8210>

I have not stepped through the source code sufficiently to explain exactly what is happening in the code you posted, but I speculate that the interaction of this Normalize instance and the BoundaryNorm make the range of values seen by the fig.colorbar different than what you expected.

In any case, if you pass norm=norm to PolyCollection, then the result looks correct:

import numpy
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.collections as mcoll
import matplotlib.colors as mcolors

numpy.random.seed(4)
N, M = 3, 3
vertices = numpy.random.random((N, M, 2))
array = numpy.random.random((1, N, 2))
# vertices = numpy.load('vertices.npy')
# array = numpy.load('array.npy')
array = numpy.ma.array(array, mask=array > .5)
# Take 2d slice out of 3D array
slice_ = array[:, :, 0:1].flatten(order='F')
fig, ax = plt.subplots()

bounds = [.1, .4, .6]
cm = mpl.colors.ListedColormap([(1.0, 0.0, 0.0), (.2, .5, .2)])
norm = mpl.colors.BoundaryNorm(bounds, cm.N)

poly = mcoll.PolyCollection(
    vertices, 
    array=slice_, 
    edgecolors='black', linewidth=.25, norm=norm)

poly.set_cmap(cm)

fig.colorbar(poly, ax=ax, orientation='vertical')
ax.add_collection(poly, autolim=True)
ax.autoscale_view()
plt.show()