MatplotLib:自定义绘图函数的颜色图选择

Question

因为有许多不同的颜色图，我想在下面的函数中设置一个参数作为颜色图的名称，它允许用户选择自己喜欢的颜色图，如"viridis"，"plasma"，"winter"，"cool"等。

def my_plot(cm_name):
# setup the normalization and the colormap
nValues = np.array(np.linspace(1, 50, 50))
normalize = mcolors.Normalize(vmin=nValues.min(), vmax=nValues.max())
colormap = cm.cm_name
# setup the colorbar
sm = cm.ScalarMappable(norm=normalize, cmap=colormap)
sm.set_array(nValues)
# below is the code for plotting

但是，当我定义如上所示的绘图函数并调用my_plot(viridis)时，会弹出一条错误消息：

NameError: name 'viridis' is not defined

谁能帮我解决这个问题？谢谢！

Answer 1

您的函数只能接受字符串作为参数，但cm.cm_name是一个属性。不能通过编写cm.'viridis'来访问cm的属性viridis。

这是getattr的完美用例，它接受一个对象和一个字符串，并返回该对象的属性，其名称与传递给getattr的字符串相同。

更改行

colormap = cm.cm_name

自

colormap = getattr(cm, cm_name)

并使用表示颜色图的字符串(不是变量(调用函数：

my_plot('viridis')

Answer 2

颜色图的名称应该是string，因此需要像my_plot("viridis")一样称呼它。如果你省略引号，python 会认为viridis是你想要传入的变量的名称，然后抱怨你没有告诉它这个变量是什么。