如何根据值给单元格着色

Question

我正在尝试使用 pyplot 来说明机器人路径规划 RL 概念中的布局。

我创建了如下布局，但无法根据单元格的值使单元格着色。我使用 np.array 生成图像和 pyplot，但 matlib 根据 std 热图对其进行着色。

我尝试了以下方法来为所有单元格着色作为开始：

ax.set_facecolor('red')

但它似乎为 ax.imshow 后面的图着色，将其渲染隐藏在插图上。

整个 .py：

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

layout = np.array([
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
 [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1],
 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
)


fig, ax = plt.subplots()
im = ax.imshow(layout)


ax.set_xticks(np.arange(len(layout[1])))
ax.set_yticks(np.arange(len(layout)))

# Rotate the tick labels and set their alignment.
plt.setp(ax.get_xticklabels(), rotation=45, ha="right",
         rotation_mode="anchor")

# Loop over data dimensions and create text annotations.
for i in range(len(layout)):
    for j in range(len(layout[1])):
        text = ax.text(j, i, layout[i, j],
                       ha="center", va="center", color="w")
        ax.set_facecolor('red')

ax.set_title("Showroom layout")
fig.tight_layout()

plt.show()

我知道 set_facecolor 现在会将整个图着色为红色，但我仍在调试。将来，它将把值为“1”的单元格着色为灰色，将“0”设置为白色。或相似。

我感谢您的帮助和意见！ :)

Answer 1

您可以定义 imshow 的工作方式。默认情况下，它使用一系列名为 viridis 的颜色，但您可以使用其他颜色。例如，灰色将为您提供介于 0 为白色和 1 为黑色之间的比例。但如果您想要灰色为 1，则需要创建自己的 cmap。见下文。它只是创建颜色的变化，白色代表最小值（此处为 0），灰色代表最大值（此处为 1）。

from matplotlib.colors import Colormap, ListedColormap
cmap = ListedColormap(["white", "gray"])
im = ax.imshow(layout, cmap=cmap)

如果您仍然想看到数字，您可能需要根据值选择文本的颜色：

# Loop over data dimensions and create text annotations.
for i in range(len(layout)):
    for j in range(len(layout[1])):
        text = ax.text(j, i, layout[i, j],
                       ha="center", va="center", color='gray' if layout[i, j] == 0 else 'white')

输出：

编辑：如果您需要两种以上的颜色并且不需要渐变而是特定的颜色，您可以使用以下颜色。它将使用每种颜色来匹配其索引（对于不太常见的颜色，您也可以使用 RGB 格式 (255,0,0)）

from matplotlib.colors import ListedColormap
cmap=ListedColormap(["red", "blue", "green", "black"])
plt.imshow(a, cmap=cmap)
# Loop over data dimensions and create text annotations.
for i in range(len(a)):
    for j in range(len(a[1])):
        text = plt.text(j, i, a[i, j],
                       ha="center", va="center", color='white')

注意：：如果您的绘图超出了颜色列表的长度，它会试图传播它，你会失去一点控制。您还可以使用 ListedColormap 的 N 参数，该参数允许您重复颜色顺序直到 N，在我的示例中，使用 N=6，颜色图将为["red", "blue", "green", "black", "red", "blue"]