plt.savefig保存EPS格式时imshow中NaN显示异常的问题

使用plt.imshow和plt.savefig保存包含NaN的二维数组图像时，调用plt.show()或保存为PNG格式时，NaN区域显示为白色，但保存为EPS文件后，NaN区域变成了黑色。以下是最小可复现代码：

import numpy as np
import matplotlib.pyplot as plt
plt.figure()
# MWE
data = np.ones([10,10]) # 生成测试数据
data[:,5] = np.nan # 将某列设为NaN

image = plt.imshow(data)
cbar = plt.colorbar(image)

plt.show()
plt.savefig("my_image.eps", format="eps")

请问这是操作失误还是Bug？如果不是Bug，有什么解决办法？

这不是操作错误，是Matplotlib处理EPS矢量格式时的渲染特性——EPS对NaN值的默认填充色，和交互式窗口、PNG位图的默认设置不一致。

可以用以下两种方法解决：

• 方法一：直接设置NaN区域的填充色
  在调用imshow后，通过色图的set_bad()方法指定NaN区域的颜色，比如设置为白色：

  import numpy as np
  import matplotlib.pyplot as plt
  plt.figure()
  data = np.ones([10,10])
  data[:,5] = np.nan
  
  image = plt.imshow(data)
  # 把NaN区域设为白色
  image.cmap.set_bad(color='white')
  cbar = plt.colorbar(image)
  
  plt.savefig("my_image.eps", format="eps")
  

• 方法二：自定义包含NaN色的色图
  如果需要更灵活的配色控制，可以基于现有色图创建自定义色图，并指定NaN颜色：

  import numpy as np
  import matplotlib.pyplot as plt
  from matplotlib.colors import ListedColormap
  
  plt.figure()
  data = np.ones([10,10])
  data[:,5] = np.nan
  
  # 基于默认的viridis色图创建自定义色图
  original_cmap = plt.cm.viridis
  new_cmap = ListedColormap(original_cmap(np.linspace(0,1,256)))
  # 指定NaN区域为白色
  new_cmap.set_bad('white')
  
  image = plt.imshow(data, cmap=new_cmap)
  cbar = plt.colorbar(image)
  
  plt.savefig("my_image.eps", format="eps")
  

内容的提问来源于stack exchange，提问作者maxxpinse