pcolormesh with masked invalid values

Question

I'm trying to plot a one-dimensional array as a pcolormesh (so the color varies along the x-axis, but is constant in the y-axis for each x). But my data has some bad values, so I'm using a masked array and a customized colormap with masked values set to blue:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import copy

a = np.array([3, 5, 10, np.inf, 5, 8])
a = np.ma.masked_where(np.isinf(a), a)
imdata = np.vstack((a, a))
myhot = copy.copy(cm.hot)
myhot.set_bad('b', 1)

fig, ax = plt.subplots()
im = ax.pcolormesh(imdata, cmap=myhot)
plt.colorbar(im)
plt.show()

It works fine if I don't have the np.inf value, but I just get a blank plot if I do. I seem to have misunderstood something about the way set_bad works because I get an additional warning:

RuntimeWarning: invalid value encountered in true_divide
  resdat /= (vmax - vmin)

What should I be doing to get the effect I want?

Answer 1

You need to mask imdata, not necessarily a:

import numpy as np
import matplotlib.pyplot as plt

a = np.array([3, 5, 10, np.inf, 5, 8])
imdata = np.ma.masked_invalid(np.atleast_2d(a))
cmap = plt.cm.hot
cmap.set_bad('b', 1)
fig, ax = plt.subplots()
im = ax.pcolormesh(imdata, cmap=cmap)

plt.colorbar(im)
plt.show()

---

If you look at imdata in an interactive session, you'll see

In [185]: imdata
Out[185]: 
masked_array(data =
 [[  3.   5.  10.  inf   5.   8.]
 [  3.   5.  10.  inf   5.   8.]],
             mask =
 False,
       fill_value = 1e+20)

Above, mask=False means that nothing is masked. If you wrap that with np.ma.masked_invalid then:

In [186]: np.ma.masked_invalid(imdata)
Out[186]: 
masked_array(data =
 [[3.0 5.0 10.0 -- 5.0 8.0]
 [3.0 5.0 10.0 -- 5.0 8.0]],
             mask =
 [[False False False  True False False]
 [False False False  True False False]],
       fill_value = 1e+20)

The problem with masking a is that np.vstack does not respect the mask. Alternatively, you could have used np.ma.vstack. Generally speaking, only functions in the np.ma namespace respect the mask.

However, you don't actually need to use vstack here; np.atleast_2d will do. vstack creates an array of shape (2, N), while np.atleast_2d creates an array of shape (1, N).

---

Another alternative is to use set_over instead of set_bad. This would allow you to avoid needing a masked array altogether:

import numpy as np
import matplotlib.pyplot as plt

a = np.array([3, 5, 10, np.inf, 5, 8])
imdata = np.atleast_2d(a)
cmap = plt.cm.hot
cmap.set_over('b')
cmap.set_under('g')
fig, ax = plt.subplots()

b = a[np.isfinite(a)]
im = ax.pcolormesh(imdata, cmap=cmap, vmin=b.min(), vmax=b.max())

plt.colorbar(im, extend='both')
plt.show()

The extend='both' in conjunction with set_over and set_under give you little colored arrows on the colorbar which indicate the color used for values beyond the colorbar's range.