如何在matplotlib中将整数映射到颜色？

Question

更新： 做了一个糟糕的修改，我选择了一个颜色贴图渐变（“nipy_spectral”在我的例子中）并制作了一个字典，将我想要的颜色映射成一个整数，这样绘图函数将绘制正确的颜色。它涉及到大量的数字。如果没有更好的解决方案，我会发布。

---

我有以下绘图功能，改编自我用于混淆矩阵的另一个绘图功能：

def plot_matrix(rm, title='Robot World', cmap=plt.cm.Blues):
    plt.imshow(rm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.tight_layout()

我有一个矩阵形式如下：

[ 1,3,1,0,1,3,4]

[ 3,1,0,1,2,3,0]

... for a few more rows.

理想情况下，我绘制此矩阵，使矩阵中的整数值对应于绘图中的颜色：['black','blue','yellow','green','red']，该绘图类似于matplotlib示例站点上的示例矩阵，用于最近的插值：

---

我想答案一定是某种定制的cmap，比如：

cmap = { 0:'k',1:'b',2:'y',3:'g',4:'r' }

在matplotlib中我该怎么做？我在看定制的彩色地图网站，它似乎比我想做的要复杂得多。。。

Answer 1

可以使用ListedColormap来执行离散颜色映射：

import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np

def plot_matrix(rm, title='Robot World', cmap=plt.cm.Blues):
    plt.imshow(rm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.tight_layout()
    plt.show()

cmap = colors.ListedColormap(['k','b','y','g','r'])
rm = np.random.randint(0,4,(5,5))
plot_matrix(rm,cmap=cmap)

，结果是：