如何在 seaborn 热图上的刻度标签上添加颜色代码？

Question

我有一个基因表达与细胞的矩阵，并希望将它们显示为热图，这本身不是问题。然而，将所有基因显示为 yticklabels 将过于混乱且视觉上没有吸引力。因此，我将每个基因注释为属于特定功能组，并希望将每个功能组表示为一种颜色，并将它们的颜色显示在热图上，与基因出现的顺序相同。澄清一下，我不想按颜色对它们进行分组，我相信你可以使用 seaborn clustermap 做到这一点。

因此，到目前为止，我有一个 pandas 数据框，其中包含基因的多索引及其各自的功能组和细胞。

我在 Stackoverflow 和 Google 上广泛搜索了答案，但没有任何运气。这是我第一次尝试这种类型的东西，所以很遗憾我不知道从哪里开始。

因此，为了简单起见，假设您有以下数据框：

import seaborn as sns
import numpy as np
import pandas as pd

data=pd.DataFrame(np.array([(0,1,2),(4,5,6),(7,8,9)]), columns=['C1','C2','C3'], index=pd.MultiIndex.from_arrays([['Gene1','Gene2','Gene3'],['A','B','A']]))

这将产生以下结果：

           C1  C2  C3
Gene1   A   0   1   2
Gene2   B   4   5   6
Gene3   A   7   8   9

现在，我可以简单地调用sns.heatmap(data)来生成热图。但是，我如何对其进行自定义，以便获得代表 A 和 B 而不是 Gene1、Gene2、Gene3 作为 yticklabels 的颜色？例如，假设 A 是蓝色，B 是绿色，我希望它显示 yticklabels（从顶部到底部）为蓝色、绿色、蓝色。

非常感谢。

Answer 1

这是一个可能的解决方案，可以在热图左侧创建一个新轴，它根据第二个 Multiindex 级别的值显示另一个热图。

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
import pandas as pd

data=pd.DataFrame(np.array([(0,1,2),(4,5,6),(7,8,9)]), 
                  columns=['C1','C2','C3'], 
                  index=pd.MultiIndex.from_arrays([['Gene1','Gene2','Gene3'],['A','B','A']]))


cats = data.index.to_frame().set_index(0)
u, inv = np.unique(cats.values, return_inverse=True)

colors = ["navy", "limegreen", "gold"]
assert(len(u) <= len(colors))

cmap = mcolors.ListedColormap(colors)
norm = mcolors.BoundaryNorm(np.arange(len(u)+1)-.5, len(u))

fig, (sax, hax) = plt.subplots(ncols=2, sharey=True,
                               gridspec_kw=dict(width_ratios=[1, data.shape[1]]))

im = sax.imshow(np.atleast_2d(inv).T, cmap=cmap, norm=norm)
hax.imshow(data.values, cmap="Greys")

sax.set_yticks(np.arange(len(cats)))
sax.set_yticklabels(cats.index)
sax.tick_params(bottom=False, labelbottom=False)

hax.set_xticks(np.arange(len(data.columns)))
hax.set_xticklabels(data.columns)

cbar = fig.colorbar(im, cax = fig.add_axes([.125, .08, .1, .04]), 
                    orientation="horizontal", ticks=np.arange(len(u)))
cbar.set_ticklabels(u)

plt.show()