How to create a discrete colormap that maps integers to colors, invariant to range of input data

Viewed 513

Let's say I have a vector containing integers from the set [1,2,3]. I would like to create a colormap in which 1 always appears as blue, 2 always appears as red, and 3 always appears as purple, regardless of the range of the input data--e.g., even if the input vector only contains 1s and 2s, I would still like those to appear as blue and red, respectively (and purple is not used in this case).

I've tried the code below:

This works as expected (data contains 1, 2 and 3):

cmap = colors.ListedColormap(["blue", "red", "purple"])
bounds = [0.5,1.5,2.5,3.5]
norm = colors.BoundaryNorm(bounds, cmap.N)

data = np.array([1,2,1,2,3])
sns.heatmap(data.reshape(-1,1), cmap=cmap, norm=norm, annot=True)

Does not work as expected (data contains only 1 and 2):

cmap = colors.ListedColormap(["blue", "red", "purple"])
bounds = [0.5,1.5,2.5,3.5]
norm = colors.BoundaryNorm(bounds, cmap.N)

data = np.array([1,2,1,2,2])
sns.heatmap(data.reshape(-1,1), cmap=cmap, norm=norm, annot=True)

In the first example, 1 appears as blue, 2 appears as red and 3 appears as purple, as desired.

In the second example, 1 appears as blue and 2 appears as purple, while red is not used.

1 Answers

Not completely sure, but I think this minimal example solves your problem. Here, I've taken an actual colormap and edited it to produce a smaller version of it. Hope it helps!

#0. Import libraries
#==============================
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import colors
import seaborn as sns
import numpy as np
#==============================

#1. Create own colormap
#======================================

#1.1. Choose the colormap you want to
#pick up colors from
source_cmap=matplotlib.cm.get_cmap('Set2')

#1.2. Choose number of colors and set a step
cols=4;step=1/float(cols - 1)

#1.3. Declare a vector to store given colors
cmap_vec=[]

#1.4. Run from 0 to 1 (limits of colormap)
#stepwise and pick up equidistant colors
#---------------------------------------
for color in np.arange(0,1.1,step):
    #store color in vector
    cmap_vec.append( source_cmap(color) )
#---------------------------------------

#1.5. Create colormap with chosen colors
    custom_cmap=\
    colors.ListedColormap([ color for color in cmap_vec ])
#====================================

#2. Basic example to plot in
#======================================
A = np.matrix('0 3; 1 2')
B=np.asarray(A)
ax=sns.heatmap(B,annot=True,cmap=custom_cmap)
plt.show()
#======================================
Related