Annotate missing values on a seaborn heatmap

Question

I am plotting a heatmap in python with the seaborn library. The dataframe contains some missing values (NaN). I wish that the heatmap cells corresponding to these fields are white (by default) and also annotated with a string NA. However, if I see it correctly, annotation does not work with missing values. Is there any hack around it?

My code:

    sns.heatmap(
        df,
        ax=ax[0, 0],
        cbar=False,
        annot=annot_df,
        fmt="",
        annot_kws={"size": annot_size, "va": "center_baseline"},
        cmap="coolwarm",
        linewidth=0.5,
        linecolor="black",
        vmin=-max_value,
        vmax=max_value,
        xticklabels=True,
        yticklabels=True,
    )

Answer 1

An idea is to draw another heatmap, with a transparent color and with only values where the original dataframe is NaN. To control the axis labels, the "real" heatmap should be drawn last. Note that the color for the NaN cells is the background color of the plot.

import seaborn as sns
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

data = np.where(np.random.rand(7, 10) < 0.2, np.nan, np.random.rand(7, 10) * 2 - 1)
df = pd.DataFrame(data)
annot_df = df.applymap(lambda f: f'{f:.1f}')
fig, ax = plt.subplots(squeeze=False)
sns.heatmap(
    np.where(df.isna(), 0, np.nan),
    ax=ax[0, 0],
    cbar=False,
    annot=np.full_like(df, "NA", dtype=object),
    fmt="",
    annot_kws={"size": 10, "va": "center_baseline", "color": "black"},
    cmap=ListedColormap(['none']),
    linewidth=0)
sns.heatmap(
    df,
    ax=ax[0, 0],
    cbar=False,
    annot=annot_df,
    fmt="",
    annot_kws={"size": 10, "va": "center_baseline"},
    cmap="coolwarm",
    linewidth=0.5,
    linecolor="black",
    vmin=-1,
    vmax=1,
    xticklabels=True,
    yticklabels=True)
plt.show()

PS: To explicitly color the 'NA' cells, e.g. cmap=ListedColormap(['yellow']) could be used.