matplotlib：為繪圖自定義顏色圖的選擇 function

Question

因為有很多不同的colormap，我想在下面的function中設置一個參數作為colormap的名稱，讓用戶可以選擇自己喜歡的colormap，例如'viridis'，'plasma'，'winter'，'酷'等等。

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

但是，當我如上所示定義繪圖 function 並調用my_plot(viridis)時，會彈出一條錯誤消息：

NameError: name 'viridis' is not defined

誰能幫我這個？ 謝謝！

Answer 1

顏色圖的名稱應該是一個string ，因此您需要像my_plot("viridis")一樣調用它。 如果省略引號，python 認為viridis是你要傳入的變量名，然后抱怨你沒有告訴它這個變量是什么。

Answer 2

您的 function 只能接受字符串作為參數，但cm.cm_name是屬性。 您不能通過編寫cm.'viridis'來訪問cm的屬性viridis 。

這是getattr的完美用例，它接受一個object和一個字符串，並返回名稱與傳遞給getattr的字符串相同的 object 的屬性。

換行

colormap = cm.cm_name

至

colormap = getattr(cm, cm_name)

並使用代表顏色圖的字符串（不是變量）調用您的 function：

my_plot('viridis')