带有自定义图例的 Seaborn 热图

Question

我有一个大型数据集，我将针对这个问题对其进行简化。我希望根据列值制作自定义图例。让我用一个例子来解释。

   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'])

我希望制作一个热图并在顶部添加一个图例以查看行是否敏感。

首先我对值进行了排序

df2 = df2.sort_values('d')

并删除了最后一列，因为它不是数字，不能在热图中。

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

在我用 sns 创建热图之后：

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

如何在哪些列敏感或不敏感的顶部添加符号？像这样？

Answer 1

您可以在热图上方的子图中绘制包含这些值的表格。

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()

哪些输出：

查看table documentation 进行进一步调整。