CODEWITHSUNDEEP
pythonpandasmatplotlib
neversaint
neversaint

Reputation: 63984

How to create predefined color range in Matplotlib heat map from a Pandas Dataframe

I have the following data frame:

import pandas as pd
Index= ['aaa', 'bbb', 'ccc', 'ddd', 'eee']
Cols = ['A', 'B', 'C', 'D']
data= [[ 1, 0.3, 2.1, 1.3],[ 2.5, 1, 1, 0.77],[ 0.0, 1, 2, 1],[ 0, 3.2, 1, 1.2],[ 10, 1, 1, 1]]
df = pd.DataFrame(data, index=Index, columns=Cols)

That looks like this:

In [25]: df
Out[25]:
        A    B    C     D
aaa   1.0  0.3  2.1  1.30
bbb   2.5  1.0  1.0  0.77
ccc   0.0  1.0  2.0  1.00
ddd   0.0  3.2  1.0  1.20
eee  10.0  1.0  1.0  1.00

What I want to do is to create a heat map with the following condition:

Ideally the color would have to be in gradation. This is my failed poor attempt:

from matplotlib import colors
cmap = colors.ListedColormap(['darkblue','blue','white','pink','red'])
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])

Which produce this plot:

enter image description here

What's the right way to do it?

Upvotes: 1

Views: 14890

Answers (1)

Marius
Marius

Reputation: 60060

To get gradiated colours you can do:

import matplotlib.pyplot as plt
# Builtin colourmap "seismic" has the blue-white-red
#   scale you want
plt.pcolor(np.array(data), cmap=plt.cm.seismic, vmin=0, vmax=2)
plt.colorbar()
plt.show()

Here I've set vmin and vmax so that they're equally spaced around the white value at 1.0, unfortunately I think this means that any values above 2.0 don't become any darker than those at 2.0. You may get better results by choosing a wider range, even if this means the scale includes negative values, e.g. vmin=-2, vmax=4.

enter image description here

Upvotes: 5

Related Questions