如何在 seaborn legend 中选择和排序值

Question

我只想在 seaborn 热图图例中包含某些值。具体来说，我有一个我不想在图例中看到的“nan”类别。

我正在尝试将医院患者的病房移动绘制为一种分类热图，不同的颜色代表不同的病房。我借用了这段代码heatmap-like plot, but for categorical variables in seaborn 来为热图配置我的输入表。空白单元格表示患者在这些日期不在医院。

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import pandas as pd

data = {'12/3': [np.nan, 'Ward_B', np.nan],
        '13/3': [np.nan, 'Ward_B', np.nan],
        '14/3': [np.nan, 'Ward_B', 'ED'],
        '15/3': ['ED', 'Ward_A', 'Ward_C'],
        '16/3': ['ED', 'Ward_A', 'Ward_C'],
        '17/3': ['Ward_A', 'Ward_A', 'Ward_C'],
        '18/3': ['Ward_A', np.nan, 'Ward_C'],
        '19/3': ['Ward_A', np.nan, 'Ward_A'],
        '20/3': [np.nan, np.nan, 'Ward_A']}

df = pd.DataFrame (data, columns = ['12/3',
                                    '13/3',
                                    '14/3',
                                    '15/3',
                                    '16/3',
                                    '17/3',
                                    '18/3',
                                    '19/3',
                                    '20/3'])

# Create dataframe of patient IDs
patient_codes_df = pd.DataFrame(['Patient_A', 'Patient_B', 'Patient_C'])
# change heading
patient_codes_df = patient_codes_df.rename(columns={0:'Patient'})
# Merge
df2 = pd.concat([patient_codes_df, df], axis=1)
# Make Patient column the index
df3 = df2.set_index('Patient')
df3

df3 是我的输入数据的样子。

这就是我绘制热图的方式

value_to_int = {j:i for i,j in enumerate(pd.unique(df3.values.ravel()))}
n = len(value_to_int)

cmap = sns.color_palette("Accent", n) # set colours

fig, ax = plt.subplots(1, 1, figsize = (6, 2), dpi=300)

mask = df3.isnull()
ax = sns.heatmap(df3.replace(value_to_int), cmap=cmap, mask=mask, linewidths=0.1, linecolor='#b5b5b5') 

ax.set_ylabel('')

# modify colorbar:
colorbar = ax.collections[0].colorbar 
r = colorbar.vmax - colorbar.vmin 
colorbar.set_ticks([colorbar.vmin + r / n * (0.5 + i) for i in range(n)])
colorbar.set_ticklabels(list(value_to_int.keys()))  
plt.xticks(rotation=90)
plt.show()

我想去掉传说中的“nan”，并重新排序，使其按照 ED、Ward_A、Ward_B、Ward_C 的合理顺序排列。

感谢您的帮助。

Answer 1

定义value_to_int 时必须删除nans。你借来的代码很好，但我想更直接的方法是在字典中手动定义你的颜色，然后在绘图时用这个替换你的 data.frame：

lvls = {'ED': 0, 'Ward_A': 1, 'Ward_B': 2, 'Ward_C': 3}
cmap = sns.color_palette("Accent", len(lvls)) # set colours

fig, ax = plt.subplots(1, 1, figsize = (6, 2), dpi=300)
sns.heatmap(df3.replace(lvls),cmap=cmap,mask=df3.isnull(),linewidths=.1,
            linecolor='#b5b5b5',ax=ax)

colorbar = ax.collections[0].colorbar 
r = colorbar.vmax - colorbar.vmin
n = len(lvls)
colorbar.set_ticks([colorbar.vmin + r / n * (0.5 + i) for i in range(n)])
colorbar.set_ticklabels(list(lvls.keys()))
plt.xticks(rotation=90)
plt.show()

Answer 2

因此，一种可行的方法是删除颜色条并手动创建我自己的图例：

# Define colours
cols = ["#ffff99", '#beaed4', '#fdc692', '#7fc97f', '#fd4396']

# Transform categorical variables into numbers for heatmap
value_to_int = {j:i for i,j in enumerate(pd.unique(df3.values.ravel()))}
n = len(value_to_int)
cmap = sns.color_palette(cols, n) # set colours

# Plot figure
fig, ax = plt.subplots(1, 1, figsize = (8, 3), dpi=300)
sns.set(font_scale=1.27, style='whitegrid')
mask = df3.isnull()
ax = sns.heatmap(df3.replace(value_to_int),
                 cmap=cmap, mask=mask, linewidths=0.1, linecolor='#b5b5b5',
                 cbar=False, # Remove the color bar legend
                 xticklabels=2)

# Tweaking figire
ax.set_ylabel('')
plt.xticks(rotation=90)

# Create a new legend
ED_patch = mpatches.Patch(color='#beaed4', label='ED')
A_patch = mpatches.Patch(color='#7fc97f', label='Ward A')
B_patch = mpatches.Patch(color='#fd4396', label='Ward B')
C_patch = mpatches.Patch(color='#3da1bf', label='Ward C')

ax.legend(handles=[ED_patch,A_patch,B_patch,C_patch],
         bbox_to_anchor=(1.22, 1),
         prop={'size': 12})