Add annotation to specific cells in heatmap

Question

I am plotting a seaborn heatmap and would like to annotate only the specific cells with custom text.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from io import StringIO

data = StringIO(u'''75,83,41,47,19
                    51,24,100,0,58
                    12,94,63,91,7
                    34,13,86,41,77''')

labels = StringIO(u'''7,8,4,,1
                    5,2,,2,8
                    1,,6,,7
                    3,1,,4,7''')

data = pd.read_csv(data, header=None)
data = data.apply(pd.to_numeric)

labels = pd.read_csv(labels, header=None)
#labels = np.ma.masked_invalid(labels)

fig, ax = plt.subplots()
sns.heatmap(data, annot=labels, ax=ax, vmin=0, vmax=100)
plt.show()

The above code generates the following heatmap:

and the commented line generates the following heatmap:

I would like to show only the non-nan (or non-zero) text on the cells. How can that be achieved?

user_9144990 · Answer

Use a string array for annot instead of a masked array:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from io import StringIO

data = StringIO(u'''75,83,41,47,19
                    51,24,100,0,58
                    12,94,63,91,7
                    34,13,86,41,77''')

labels = StringIO(u'''7,8,4,,1
                    5,2,,2,8
                    1,,6,,7
                    3,1,,4,7''')

data = pd.read_csv(data, header=None)
data = data.apply(pd.to_numeric)

labels = pd.read_csv(labels, header=None)
#labels = np.ma.masked_invalid(labels)

# Convert everything to strings:
annotations = labels.astype(str)
annotations[np.isnan(labels)] = ""

fig, ax = plt.subplots()
sns.heatmap(data, annot=annotations, fmt="s", ax=ax, vmin=0, vmax=100)
plt.show()

user_10659910 · Answer

To complement the answer by @mrzo, you can use na_filter=False in read_csv() to store nans as empty strings and use pandas.DataFrame.astype() to convert to strings in place:

# ...
labels = pd.read_csv(labels, header=None, na_filter=False).astype(str)
sns.heatmap(data, annot=labels, fmt='s', ax=ax, vmin=0, vmax=100)

user_4041902 · Answer

Just going to add this as it has taken me some time to work out how to do something similar programmatically for a slightly different application: I wanted to suppress 0-values from the annotation, but because the values were arising as the result of a crosstab operation I couldn't use William Miller's nice approach without writing the crosstab out and then reading it back in which seemed... inelegant.

There may be a yet more elegant way to do this, but for me running it through numpy was ridiculously fast and quite easy.

import numpy as np
import pandas as pd
import seaborn as sns

from io import StringIO

data = StringIO(u'''75,83,41,47,19
                    51,24,100,0,58
                    12,94,63,91,7
                    34,13,86,41,77''')

data = pd.read_csv(data, header=None)
data = data.apply(pd.to_numeric)

# For more complex functions you could write a def instead
# of using this simple lambda function
an = np.vectorize(lambda x: '' if x<50 else str(round(x,-1)))(data.to_numpy())

sns.heatmap(
    data=data.to_numpy(), # Note this is now numpy too
    cmap='BuPu',
    annot=an,   # The matching ndarray of annotations
    fmt = '',   # Formats annotations as strings (i.e. no formatting)
    cbar=False, # Seems overkill if you've got annotations
    vmin=0, 
    vmax=data.max().max()
)

This can make life a little more difficult in terms of labelling axes, though it's straightforward enough: ax.set_xticklabels(df.columns.values). And if you had axislabels in, say, the first column then you'd need to use iloc (data.iloc[:,1:]) in your to_numpy call, but combined with a custom colormap (e.g. 0==white) you can create heatmaps that are a lot easier to look at.

Obviously the crude rounding is confusing (why does 80 have different shades?) but you get the point:

Add annotation to specific cells in heatmap

3 Answers