ฉันจะพล็อตค่า NaN เป็นสีพิเศษด้วย imshow ใน matplotlib ได้อย่างไร

ฉันพยายามใช้ imshow ใน matplotlib เพื่อพล็อตข้อมูลเป็นแผนที่ความร้อน แต่ค่าบางค่าเป็น NaN ฉันต้องการให้ NaN แสดงผลเป็นสีพิเศษที่ไม่พบใน colormap

ตัวอย่าง:

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

ภาพที่ได้จะเป็นสีฟ้าทั้งหมดโดยไม่คาดคิด (สีต่ำสุดในเจ็ท Colormap) อย่างไรก็ตามถ้าฉันวางแผนเช่นนี้:

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

- แล้วฉันจะได้อะไรที่ดีขึ้น แต่ค่า NaN ถูกวาดเป็นสีเดียวกับ vmin ... มีวิธีที่สง่างามไหมที่ฉันสามารถตั้งค่า NaN ให้วาดด้วยสีพิเศษ (เช่นเทาหรือโปร่งใส)

ด้วย Matplotlib เวอร์ชันใหม่กว่าคุณไม่จำเป็นต้องใช้อาร์เรย์มาสก์อีกต่อไป

ตัวอย่างเช่นให้สร้างอาร์เรย์โดยทุกค่าที่ 7 เป็น NaN:

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

เราสามารถแก้ไข colormap ปัจจุบันและพล็อตอาร์เรย์ด้วยบรรทัดต่อไปนี้:

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

อืมดูเหมือนว่าฉันสามารถใช้อาร์เรย์มาสก์เพื่อทำสิ่งนี้:

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)

แค่นี้ก็เพียงพอแล้วแม้ว่าฉันจะยังเปิดรับข้อเสนอแนะ :]