3D表面图中的数组的布尔掩码破坏了色彩映射

Question

我尝试3D绘制函数fun并使用colormap来显示函数值的级别。我想在非sqaured区域上绘制此函数，因此我使用布尔掩码将np.nan设置为meshgrid中的某些值。但我得到了

RuntimeWarning: invalid value encountered in less
  cbook._putmask(xa, xa < 0.0, -1)

每当我添加布尔掩码时似乎该错误是由于np.nan无法在colormap中进行比较。但我无法找到解决这个问题的方法。

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

%matplotlib inline

fun = lambda x: np.sin(x[0])*np.exp(1-np.cos(x[1]))**2 + np.cos(x[1])*np.exp(1-np.sin(x[0]))**2 + (x[0]-x[1])**2

fig = plt.figure(figsize=(8, 5))
ax = fig.gca(projection='3d')
x = np.arange(-6, 6, 3e-2)
y = np.arange(-6, 6, 3e-2)

# A constraint on x and y
x, y = np.meshgrid(x, y)
r2 = (x+5)**2 + (y+5)**2
scope = r2 < 25
# Mask is the cause of the problem
x[scope] = np.nan
y[scope] = np.nan
z = fun(np.array([x, y]))

surf=ax.plot_surface(x, y, z, cmap=cm.jet)
ax.contourf(x, y, z, offset=-120, cmap=cm.jet)
fig.colorbar(surf)
ax.view_init(elev=30, azim=60)

Answer 1

您无法修复运行时警告。这是基于阵列中存在nan值的事实的警告。

为了仍然获得颜色编码的曲面图，您可以使用matplotlib.colors.Normalize实例告诉曲面图使用哪种颜色。

请参阅下面的完整代码：

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
import matplotlib.colors

fun = lambda x: np.sin(x[0])*np.exp(1-np.cos(x[1]))**2 + np.cos(x[1])*np.exp(1-np.sin(x[0]))**2 + (x[0]-x[1])**2

fig = plt.figure(figsize=(8, 5))
ax = fig.gca(projection='3d')
x = np.arange(-6, 6, 3e-2)
y = np.arange(-6, 6, 3e-2)

# A constraint on x and y
x, y = np.meshgrid(x, y)
r2 = (x+5)**2 + (y+5)**2
scope = r2 < 25
# Mask is the cause of the problem
x[scope] = np.nan
y[scope] = np.nan
z = fun(np.array([x, y]))

norm = matplotlib.colors.Normalize(vmin=-120, vmax=120)
cm.jet.set_under((0,0,0,0))
ax.contourf(x, y, z, offset=-120, cmap=cm.jet, norm=norm)
surf=ax.plot_surface(x, y, z, cmap=cm.jet, norm=norm)
fig.colorbar(surf)
#ax.view_init(elev=30, azim=60)
plt.show()