Search code examples
pythonpandasmatplotlibseaborn

Customize text annotations of heatmaps in a Seaborn FacetGrid to be in EngFormatter style


While I am trying to produce a heatmap compare the performance of different user segments across time periods, I cannot customize the annotation of the heatmaps in EngFormatter style (i.e. 1233 = 1k, 10000 = 10k, 2000000 = 2M, etc.). I could not find where is the location of the annotation text element under each heatmap within the FacetGrid.

How to customize the unit of Facegrid Annotation?

My effort:

import matplotlib.pyplot as plt
import seaborn as sns

def draw_heatmap(*args, **kwargs):
    data = kwargs.pop('data')
    d = data.pivot(index=args[1], columns=args[0], values=args[2])
    sns.heatmap(d, **kwargs)

fg = sns.FacetGrid(df, col='time_interval',height=6)
fg.map_dataframe(
    draw_heatmap,
    'user_action', 'segment', 'users',
    cbar=True, cmap="Blues", square = True, annot=True,fmt='.2g'
)

from matplotlib.ticker import EngFormatter
for t in fg.texts: t.set_text(EngFormatter(t.get_text()))

Current output:

enter image description here

Desired output:

A visualization like above but where the displayed text annotations of heatmaps are in EngFormatter style

Sample dataset:

index,time_interval,user_action,segment,users
0,time_2,,Super Heavy,1233
1,time_1,Click Use,Medium,10000
2,time_1,Click View Detail,Light,2000000
3,time_2,Click View Detail,Medium,3999
4,time_1,Click See All,Medium,14542

Session info:

polars==0.17.3
pandas==1.5.3
seaborn==0.12.2

Solution

  • You were very close to getting it right, but put the iteration and EngFormatter alteration of the texts attribute of the Axes object within the draw_heatmap function. You also need to capture the Axes of the seaborn heatmap plots themselves to modify their annotations (i.e., not the texts attribute of the FacetGrid, but of the heatmaps within the grid).

    E.g., using your provided sample data (slightly modified for extended demonstration):

    import itertools as itl
    import random
    
    import matplotlib.pyplot as plt
    import pandas as pd
    import seaborn as sns
    
    from matplotlib.ticker import EngFormatter
    
    times = ["time_1", "time_2"]
    user_actions = ["Click See All", "Click Use", "Click View Detail"]
    segments = ["Light", "Medium", "Heavy"]
    
    all_combos = list(map(list, list(itl.product(times, user_actions, segments))))
    
    cols = ["time_interval", "user_action", "segment", "users"]
    
    data = [c + [random.randint(1, int(2e6))] for c in all_combos]
    
    df = pd.DataFrame(data, columns=cols)
    
    
    def draw_heatmap(*args, **kwargs):
        data = kwargs.pop("data")
        d = data.pivot(index=args[1], columns=args[0], values=args[2])
        ax = sns.heatmap(d, **kwargs)
    
        # Make annotations of heatmaps in EngFormatter style
        for text in ax.texts:
            text.set_text(
                EngFormatter()(float(text.get_text()))
            )
    
    
    fg = sns.FacetGrid(df, col="time_interval", height=6)
    fg.map_dataframe(
        draw_heatmap,
        "user_action",
        "segment",
        "users",
        cbar=True,
        cmap="Blues",
        square=True,
        annot=True,
        fmt=".2g",
    )
    
    plt.show()
    

    produces: EngFormatter text annotated seaborn facet grid heatmaps

    where the randomly generated data ('df') shown above was:

       time_interval        user_action segment    users
    0         time_1      Click See All   Light   744670
    1         time_1      Click See All  Medium   150098
    2         time_1      Click See All   Heavy   588170
    3         time_1          Click Use   Light  1811019
    4         time_1          Click Use  Medium   495390
    5         time_1          Click Use   Heavy  1312397
    6         time_1  Click View Detail   Light  1598512
    7         time_1  Click View Detail  Medium  1763711
    8         time_1  Click View Detail   Heavy   930745
    9         time_2      Click See All   Light  1832730
    10        time_2      Click See All  Medium  1868217
    11        time_2      Click See All   Heavy     4131
    12        time_2          Click Use   Light  1163743
    13        time_2          Click Use  Medium   600851
    14        time_2          Click Use   Heavy  1019326
    15        time_2  Click View Detail   Light  1527671
    16        time_2  Click View Detail  Medium   559187
    17        time_2  Click View Detail   Heavy  1943017