Custom colors in matplotlib when using matshow

Question

Is there an easy way to indicate a specific color for each element of a given matrix when using matplotlib. For example, assume we want to show 'x' as follow with three specific colors: red, black, and, white:

However, the only option I found out is using "cmap" which doesn't directly give you the option to "directly" specify the colors.

fig = plt.figure()
ax = fig.add_subplot(111)
x= [[0,0,0,0,0,0],[0,0,0,0,0,0], [0,1,1,2,1,1], [0,0,0,0,0,1], [0,1,1,1,1,1]]    
cax = ax.matshow(x,cmap=plt.cm.gray_r )
plt.show()

My question: how should I change my code to show the above red/black/white grid? [e.g 0 means black, 1 means white, and 2 means red] and in general how we can do it for a larger list of colors? like 10-15 colors.

In addition, how to assign to a certain element in the matix a certain color? for example in above, x[i][j] == 0 then color ='black' or x[i][j] == 2 then color ='red'

Thanks.

Answer 1

You can create your own color maps:

from matplotlib.colors import ListedColormap

cmap = ListedColormap(['k', 'w', 'r'])
cax = ax.matshow(x,cmap=cmap)

If you want to specify 10-15 colors you may run out of single-letter colors. In this case you can specify RGB triplets (e.g. ListedColormap([[0, 0, 0], [1, 1, 1], [1, 0, 0]])) or various other color formats. Alternatively, use one of the pre-defined discrete ("qualitative") color maps listed here.

If the values in the matrix are not consecutive integers you can transform them before plotting.

import numpy as np

x = np.array([[0,0,0,0,0,0],[0,77,0,0,22,0], [0,1,1,2,1,1], [0,0,14,0,0,1], [0,1,1,1,1,1]])
u, i = np.unique(x, return_inverse=True)
y = i.reshape(x.shape)
# array([[0, 0, 0, 0, 0, 0],
#        [0, 5, 0, 0, 4, 0],
#        [0, 1, 1, 2, 1, 1],
#        [0, 0, 3, 0, 0, 1],
#        [0, 1, 1, 1, 1, 1]])