Seaborn heatmap with a custom legend

Question

I have a large dataset which I will simplify for this question. I wish to make a custom legend based on the column value. Let me explain with an example.

   df2 = DataFrame(np.array([[1.0, 2.0, 3.0, 'sensitve'], [4.0, 5.0, 6.0, 'sensitive'], [7.0, 8.0, 9.0, 'not sensitive']]),
                   columns=['a', 'b', 'c', 'd'], index = ['A', 'B', 'C'])

I wish to make a heatmap and have a legend on top to see if row is sensitive or not.

First I sorted the values

df2 = df2.sort_values('d')

and deleted the last column, because it isn't a number and can't be in a heat map.

df2 = df2.iloc[:, :-1].astype(float).T

After I created a heatmap with sns:

sns.heatmap(df2, yticklabels = False)#, cmap="viridis")

How do I make a notation up the top of which columns are sensitive or not? Something like this?

Answer 1

You could plot a table with those values in a subplot above the heatmap.

df2 = pd.DataFrame(np.array([[1.0, 2.0, 3.0, 'sensitve'], [4.0, 5.0, 6.0, 'sensitive'], [7.0, 8.0, 9.0, 'not sensitive']]),
               columns=['a', 'b', 'c', 'd'], index = ['A', 'B', 'C'])
df2 = df2.sort_values('d')

d_col = df2['d'].T
df2 = df2.iloc[:, :-1].astype(float).T

# set subplot height ratios to avoid large empty space above the table
fig, (ax_table,ax_heatmap) = plt.subplots(2,1, figsize = [8,8], gridspec_kw = {'height_ratios':[1,9]})
sns.heatmap(df2, linewidth = 1, ax = ax_heatmap)
# only using 2 unique d_col values and setting colWidths to fit your drawing
ax_table.table([d_col.unique()], colWidths = [1/3,2/3], cellLoc = 'center')
ax_table.axis('off')

# setting table size to fit heatmap following https://stackoverflow.com/questions/66825885/centering-a-table-with-a-heatmap
bbox_heatmap = ax_heatmap.get_position()
bbox_table = ax_table.get_position()    
left = bbox_heatmap.x0
bottom = bbox_table.y0
width = bbox_heatmap.x0 + (bbox_heatmap.width * 0.8)
height = bbox_table.height * 1.2    
ax_table.set_position([left, bottom, width, height])

plt.show()

Which outputs:

Check out the table documentation for further adjustments.