Как выбрать и упорядочить значения в seaborn legend

Question

Я хотел бы включить в легенду морской карты только определенные значения. В частности, у меня есть категория «нан», которую я не хочу видеть в легенде.

Я пытаюсь изобразить движения палат для пациентов в больнице как своего рода категориальную тепловую карту с разными цветами, представляющими разные подопечные. Я позаимствовал из этого кода график, подобный тепловой карте, но для категориальных переменных в 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

Итак, один из подходов, который сработал, заключался в том, чтобы удалить цветовую полосу и создать свою собственную легенду вручную:

# 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})

Answer 2

Вы должны отбросить nans, когда определяете value_to_int. Код, который вы заимствовали, хорош, но я думаю, что более простой способ - определить ваши цвета вручную в словаре, а затем заменить свой 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()