Creating heatmaps with matplotlib in python

Posted on | In

Creating (categorical) heatmaps with matplotlib consists of 2 steps: painting cells with colors and annotating it.

Packages

The package in need is mainly matplotlib.

1
2
3
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

Using numpy or pandas would also make things easier.

Painting cells

Assuming we already have an numpy array named fusMat (i.e. a fusion matrix). The name of columns (as well as rows) in a list is classes. To paint the cells, do this

1
2
fig, ax = plt.subplots()
im = ax.imshow(fusMat)

and we will get something like: To change the coloring scheme, use parameter cmap in .imshow() method:

1
im = ax.imshow(fusMat, cmap=plt.cm.Reds)

Annotating

Clearly the axises in the figure above is not what we want. So we need to change these “tick labels”:

1
2
3
4
ax.set_xticks(np.arange(len(classes)))
ax.set_yticks(np.arange(len(classes)))
ax.set_xticklabels(classes)
ax.set_yticklabels(classes)

We would also like to specify what x- and y-axis stands for:

1
2
ax.set_xlabel("pred")
ax.set_ylabel("gold")

A color bar would be nice:

1
2
cbar = ax.figure.colorbar(im, ax=ax)
cbar.ax.set_ylabel(ylabel="# of pred", rotation=-90, va="bottom")

We want the numbers of each cell be displayed, in white on dark cells, and black on light cells:

1
2
3
4
5
6
7
8
for i in range(0,len(classes)):
    for j in range(0,len(classes)):
        if fusMat[j][i] > (fusMat.max()+fusMat.min())/2:
            text = ax.text(i,j, fusMat[j][i],
                           ha="center", va="center", color="w")
        else:
            text = ax.text(i,j, fusMat[j][i],
                           ha="center", va="center", color="k")

To add title, use

1
ax.set_title("Fusion matrix")

Finally, we get somthing like this:

Saving to file

To save the fig to a file:

1
2
fig.tight_layout()
plt.savefig(filename)

References