¿Cómo puedo trazar los valores de NaN como un color especial con imshow en matplotlib?

Estoy tratando de usar imshow en matplotlib para trazar datos como un mapa de calor, pero algunos de los valores son NaN. Me gustaría que los NaN se representen como un color especial que no se encuentra en el mapa de colores.

ejemplo:

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()

La imagen resultante es inesperadamente toda azul (el color más bajo en el mapa de colores del jet). Sin embargo, si hago el trazado así:

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

- luego obtengo algo mejor, pero los valores de NaN se dibujan del mismo color que vmin ... ¿Hay alguna manera elegante de configurar los NaN para que se dibujen con un color especial (por ejemplo: gris o transparente)?

Con las versiones más recientes de Matplotlib, ya no es necesario usar una matriz enmascarada.

Por ejemplo, generemos una matriz donde cada séptimo valor es un NaN:

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

Podemos modificar el mapa de colores actual y trazar la matriz con las siguientes líneas:

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

Hrm, parece que puedo usar una matriz enmascarada para hacer esto:

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)

Esto debería ser suficiente, aunque todavía estoy abierto a sugerencias. :]