python - 使用多索引数据框时如何在 seaborn 热图中自定义 y 标签？

Question

我有一个包含多索引行的 DataFrame，我想创建一个没有重复行标签的热图，就像它出现在 pandas DataFrame 中一样。这是复制我的问题的代码:

import pandas as pd
from matplotlib import pyplot as plt
import random
import seaborn as sns
%matplotlib inline

df = pd.DataFrame({'Occupation':['Economist','Economist','Economist','Engineer','Engineer','Engineer',
                           'Data Scientist','Data Scientist','Data Scientist'],
             'Sex':['Female','Male','Both']*3, 'UK':random.sample(range(-10,10),9),
             'US':random.sample(range(-10,10),9),'Brazil':random.sample(range(-10,10),9)})

df = df.set_index(['Occupation','Sex'])

df

sns.heatmap(df, annot=True, fmt="",cmap="YlGnBu")

除了消除重复之外，我还想自定义一些 y 标签，因为这种原始形式对我来说看起来不太好。

这可能吗？

Answer 1

据我所知，在 seaborn 中没有快速简便的方法来做到这一点，但希望有人能纠正我。您可以通过将 ytick_labels 重置为索引级别 1 的值来手动执行此操作。然后，您可以遍历索引的 0 级，并在正确的位置将 text 元素添加到您的可视化中:

from collections import OrderedDict

ax = sns.heatmap(df, annot=True, cmap="YlGnBu")

ylabel_mapping = OrderedDict()
for occupation, sex in df.index:
    ylabel_mapping.setdefault(occupation, [])
    ylabel_mapping[occupation].append(sex)
    
hline = []
new_ylabels = []
for occupation, sex_list in ylabel_mapping.items():
    sex_list[0] = "{} - {}".format(occupation, sex_list[0])
    new_ylabels.extend(sex_list)
    
    if hline:
        hline.append(len(sex_list) + hline[-1])
    else:
        hline.append(len(sex_list))


ax.hlines(hline, xmin=-1, xmax=4, color="white", linewidth=5)
ax.set_yticklabels(new_ylabels)

另一种方法涉及使用数据框样式。这导致了一个 super 简单的语法，但你确实失去了 colobar。这使您的索引和列显示与数据框完全相同。请注意，您需要在笔记本或可以呈现 html 的地方工作以查看输出:

df.style.background_gradient(cmap="YlGnBu", vmin=-10, vmax=10)