Search code examples
pythonpandasseabornheatmapplot-annotations

How to add custom annotations with uncertainty to a heatmap


I am attempting to visualize some data as a table, where the boxes of each table element are colored according to their value, the numerical value is also displayed, and the uncertainty on each element is shown. I can achieve 2 out of these 3 things using pandas.pivot_table and sns.heatmap, but cannot seem to include the uncertainty on each table element as part of the annotation. In the example code snippet:

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

df = pd.DataFrame({"A": ["foo", "foo", "foo", "foo", "foo",
                         "bar", "bar", "bar", "bar"],
                   "B": ["one", "one", "one", "two", "two",
                         "one", "one", "two", "two"],
                   "C": ["small", "large", "large", "small",
                         "small", "large", "small", "small",
                         "large"],
                   "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
                   "E": [2, 4, 5, 5, 6, 6, 8, 9, 9]})

table = pd.pivot_table(df, values='D', index=['A', 'B'],
                       columns=['C'], aggfunc=np.sum, fill_value=0)

sns.heatmap(table,annot=True)

we produce a table like so:

seaborn heatmap of a pandas pivot_table

However, imagine that the entries "E" represented the uncertainty on elements "D". Is there any way these can be displayed on the table, as "E"[i]+/-"D"[i]? I tried using a custom annotation grid, but this requires a numpy array and so string formatting each element didn't work for this.


Solution

  • You can pass a DataFrame with the formatted strings to sns.heatmap:

    table = pd.pivot_table(df, values=['D', 'E'], index=['A', 'B'],
                           columns=['C'], aggfunc=np.sum, fill_value=0)
    
    sns.heatmap(table['D'],
                annot=table['D'].astype(str)+'±'+table['E'].astype(str),
                fmt='')
    

    enter image description here