How to map integers to colors in matplotlib?

Question

Update: Did a lousy hack where I choose a colormap gradient ('nipy_spectral' in my case) and make a dictionary that maps the color I want to an integer such that the plotting function will plot the correct color. It involves a lot of futzing around with the numbers. I'll post as solution if nothing better comes along.

I have the following plotting function adapted from another plotting function I used for a confusion matrix:

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

And I have a matrix of the following form:

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

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

... for a few more rows.

Ideally, I plot this matrix such that the integer values within the matrix correspond to the colors: ['black','blue','yellow','green','red'] in a plot that looks something like the example matrix on matplotlib's example site for interpolation nearest:

I think the answer must be some kind of custom cmap, like:

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

How would I do this in matplotlib? I am looking at the custom colormap site, and it seems a whole lot more involved than what I am trying to do...

Also this https://stackoverflow.com/questions/14777066/matplotlib-discrete-colorbar/14777392#14777392. — Reti43, Apr 02 '16 at 19:59
From the above link the mapping is: `cmap = plt.cm.get_cmap('PiYG', 11)` `plt.plot(range(5), c=cmap(0))` — user3226167, Jun 11 '19 at 02:58

score 11 · Accepted Answer · edited Apr 01 '21 at 03:38

You can use a ListedColormap to do a discrete colormap:

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)

, the result is this:

How to map integers to colors in matplotlib?

1 Answers1

Linked