displaying a colored 2d array in matplotlib in Python

Question

I'd like to plot a 2-d matrix from numpy as a colored matrix in Matplotlib. I have the following 9-by-9 array:

my_array = diag(ones(9))

# plot the array
pcolor(my_array)

I'd like to set the first three elements of the diagonal to be a certain color, the next three to be a different color, and the last three a different color. I'd like to specify the color by a hex code string, like "#FF8C00". How can I do this?

Also, how can I set the color of 0-valued elements for pcolor?

Answer 1

To have the elements be different colors, assign them different values:

my_array = diag([1,1,1,2,2,2,3,3,3])

To specify the colors, try:

from matplotlib.colors import ListedColormap, NoNorm
cmap = ListedColormap(['#E0E0E0', '#FF8C00', '#8c00FF', '#00FF8C'])
pcolor(my_array,cmap=cmap,norm=NoNorm())

The norm=NoNorm() argument avoids any scaling of the matrix values, so that 0 gets the first color in the list, 1 the second, etc.