如何为seaborn heatmap添加行分组颜色标识（非聚类）

你可以通过Matplotlib的子图布局，在热图左侧添加颜色侧边栏来实现分组标识，同时完全保留热图的原始行顺序（不执行聚类操作）。以下是完整实现代码：

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

# 生成示例数据
data = pd.DataFrame(np.random.randn(10, 3), columns=list('ABC'))
data['group'] = pd.Series(random.choices(['group_1', 'group_2', 'group_3'], weights=[1,1,1], k=len(data)))

# 定义分组-颜色映射字典
lut = {'group_1': 'red', 'group_2': 'blue', 'group_3': 'green'}
row_colors = data['group'].map(lut)

# 分离数值数据和分组信息
heatmap_data = data.drop(['group'], axis=1)

# 创建子图布局：左侧颜色栏（窄幅）+ 右侧主热图
fig, (ax_color, ax_heatmap) = plt.subplots(ncols=2, figsize=(8, 6), gridspec_kw={'width_ratios': [0.05, 0.95]})

# 绘制左侧分组颜色栏
ax_color.imshow(row_colors.values.reshape(-1,1), aspect='auto', interpolation='nearest')
# 隐藏颜色栏的刻度与边框，保持简洁
ax_color.set_xticks([])
ax_color.set_yticks([])
for spine in ax_color.spines.values():
    spine.set_visible(False)

# 绘制主热图（不聚类，保留原始顺序）
sns.heatmap(heatmap_data, ax=ax_heatmap, cbar=False)

# 添加自定义分组图例
from matplotlib.patches import Patch
legend_elements = [Patch(facecolor=color, label=group) for group, color in lut.items()]
plt.legend(handles=legend_elements, bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0)

plt.tight_layout()
plt.show()

核心逻辑说明 #

• 子图拆分：通过gridspec_kw控制颜色栏和热图的宽度比例，让颜色栏仅作为标识存在
• 颜色栏绘制：用imshow将分组颜色转换成垂直纯色块，模拟clustermap的行颜色效果
• 热图保留原始顺序：直接使用sns.heatmap绘制原始数据，不启用任何聚类参数
• 自定义图例：通过Patch生成对应分组的颜色标记，确保图例与颜色栏一一对应

内容的提问来源于stack exchange，提问作者Shuman_tov