How to color cells in a plt plot based on value?-CodePudding

Im trying to use pylot to illustrate a layout in a robotic path planning RL concept. I've created the layout as below, but I cannot get the cells to color based on the value of the cell. I've used an np.array to generate the image and pyplot, but matlib colors it based on a std heatmap.

I have tried the following just to color all cells as a start:

ax.set_facecolor('red')

But it seem to color the plot behind the ax.imshow, rendering it hidden on the illustration.

The entire .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()

I know the set_facecolor right now will color the entire plot red, but im still debugging. In the future it will color cells with value "1" gray and "0" white. Or alike.

I appreciate your help and input! :)

CodePudding user response：

You can define the way the imshow works. By default, it uses a range of colors called viridis but you can use others. For instance, gray will give you a scale between white for 0 and black for 1. But if you want gray for 1, you need to create your own cmap. See below. It just creates a variation of color that is white for the min value (here 0) and gray for the max value (here 1).

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

if you want to still see the numbers, you may want to chose the color of the text based on the value:

# 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')

Output:

EDIT: If you need more than 2 colors and don't wan't a gradient but rather specific colors you can use the below. It will use each color to match its index (you can use RGB format too (255,0,0) for less common colors)

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')

Note:: if the numbers in your plot exceed the length of the color list, it will try to spread it and you will lose a bit the control. You could also use the N parameter of the ListedColormap that allows you to repeat the color order until N, in my example, with N=6, the colormap would be ["red", "blue", "green", "black", "red", "blue"]