如何在色条中忽略NaN？

Question

我有一个3D表面图，并且有一个颜色条是通过以下代码生成的：

surf = ax.plot_surface(xv, yv, zv, ...)
cb = fig.colorbar(surf)

当它正常工作时，效果是这样的：

但问题是，有些值可能是NaN（表示“不是一个数字”），这样颜色条就无法生成，像这样：

C:\Users\Sam\Anaconda\lib\site-packages\matplotlib\colorbar.py:581: RuntimeWarning: invalid value encountered in greater
  inrange = (ticks > -0.001) & (ticks < 1.001)
C:\Users\Sam\Anaconda\lib\site-packages\matplotlib\colorbar.py:581: RuntimeWarning: invalid value encountered in less
  inrange = (ticks > -0.001) & (ticks < 1.001)
C:\Users\Sam\Anaconda\lib\site-packages\matplotlib\colors.py:576: RuntimeWarning: invalid value encountered in less
  cbook._putmask(xa, xa < 0.0, -1)

我可以尝试把zv中的NaN值替换成0（zv[isnan(zv)] = 0），但是这样生成的图就会出现垂直的悬崖，看起来很混乱，还会遮挡一些特征，并且会影响颜色条的显示。

我想要的是一种方法，让颜色条忽略NaN值。为此，我可以在创建颜色条后手动设置它（设置为amin(zv[~isnan(zv)])和amax(zv[~isnan(zv)])），但我不知道该怎么做。

有没有什么建议可以在计算颜色条和给3D表面上色时忽略NaN值呢？

Answer 1

你应该在进行surf之前，把zv中的NaN（不是数字）排除掉，可以这样做：

surf = ax.plot_surface(xv, yv, zv[~isnan(zv)], ...)
cb = fig.colorbar(surf)

这样做应该能完全避免这个问题，如果我理解得没错的话。当然，我可能也会搞错。祝你好运。

Answer 2

在调用 plot_surface 时，设置 vmin 和 vmax 参数。这些值会影响颜色条所映射的最小值和最大值。

import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as axes3d

def peaks(x, y):
    return x*np.sin(y)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
X, Y = np.meshgrid(np.linspace(-3, 3, 100), np.linspace(-3, 3, 100))

Z = peaks(X, Y)
Z[X+Y>3] = np.nan

surf = ax.plot_surface(X, Y, Z, alpha=0.3, cmap=plt.cm.jet,
                       vmin=np.nanmin(Z), vmax=np.nanmax(Z))
cbar = plt.colorbar(surf)

plt.show()