Como posso plotar valores NaN como uma cor especial com imshow em matplotlib?

Estou tentando usar imshow em matplotlib para plotar dados como um mapa de calor, mas alguns dos valores são NaNs. Gostaria que os NaNs fossem renderizados como uma cor especial não encontrada no mapa de cores.

exemplo:

import numpy as np
import matplotlib.pyplot as plt
f = plt.figure()
ax = f.add_subplot(111)
a = np.arange(25).reshape((5,5)).astype(float)
a[3,:] = np.nan
ax.imshow(a, interpolation='nearest')
f.canvas.draw()

A imagem resultante é inesperadamente toda azul (a cor mais baixa no mapa de cores do jato). No entanto, se eu fizer a plotagem assim:

ax.imshow(a, interpolation='nearest', vmin=0, vmax=24)

--então consigo algo melhor, mas os valores NaN são desenhados da mesma cor que vmin ... Existe uma maneira elegante de definir NaNs para serem desenhados com uma cor especial (por exemplo: cinza ou transparente)?

Com as versões mais recentes do Matplotlib, não é mais necessário usar um array mascarado.

Por exemplo, vamos gerar uma matriz onde cada sétimo valor é um NaN:

arr = np.arange(100, dtype=float).reshape(10, 10)
arr[~(arr % 7).astype(bool)] = np.nan

Podemos modificar o mapa de cores atual e plotar a matriz com as seguintes linhas:

current_cmap = matplotlib.cm.get_cmap()
current_cmap.set_bad(color='red')
plt.imshow(arr)

Hrm, parece que posso usar uma matriz mascarada para fazer isso:

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

Isso deve bastar, embora eu ainda esteja aberto a sugestões. :]