Customizing colors in matplotlib - heatmap

Question

How can I specify colors in heatmap. In this example, the data are uniquely one of 4 values {0,1,2,3}

Index= ['aaa', 'bbb', 'ccc', 'ddd', 'eee']
Cols = ['A', 'B', 'C', 'D']

data= [[ 0, 3, 1, 1],[ 0, 1, 1, 1],[ 0, 1, 2, 1],[ 0, 2, 1, 2],[ 0, 1, 1, 1]]
print data
df = pd.DataFrame(data, index=Index, columns=Cols)
heatmap = plt.pcolor(np.array(data))
plt.colorbar(heatmap)
plt.show()

How can I specifiy those colors in a way to represent colors= {0:'green',1:'red',2:'black',3:'yellow'}

Answer 1

Create custom colormap and set ticks to your integers

from matplotlib import colors
cmap = colors.ListedColormap(['green','red','black','yellow'])
bounds=[-0.5, 0.5, 1.5, 2.5, 3.5]
norm = colors.BoundaryNorm(bounds, cmap.N)
heatmap = plt.pcolor(np.array(data), cmap=cmap, norm=norm)
plt.colorbar(heatmap, ticks=[0, 1, 2, 3])

Is this what you want? Notice, that your data are displayed "upside down".

Answer 2

I modified this code to show 3 red / yellow / green states of 9 nodes

import matplotlib.pyplot as plt
from matplotlib.colors 
import LinearSegmentedColormap
colors = [(1, 0, 0), (1, 1, 0), (0, 1, 0)]  # Red, yellow, green
n_bins = [3]  # Discretizes the interpolation into bins 
cmap_name = 'my_list' 
cm = LinearSegmentedColormap.from_list(cmap_name, colors, N=3)
threshold = 3  # max value
data = [[1, 1, 2], [1, 1, 3], [1, 1, 2]]
img = plt.imshow(data, interpolation='nearest', vmax=threshold, cmap=cm)
plt.show()