如何使用Seaborn基于ROTATION列绘制带EID标签的分类热力图

原有代码错误点 #

• sns.load_dataset("df") 是seaborn加载内置示例数据集的接口，无法读取你自己构造的df变量，属于无效调用
• pivot函数参数使用完全错误，你当前的数据集已经是EID为行、ROTATION1-5为列的结构，无需再做透视，仅需将EID设置为行索引即可
• 热力图仅支持数值型数据映射颜色，你当前数据混合了字符串分类和数值，需要先做值到数字的统一映射
• 原代码中变量名前后不一致，定义了in_range后续却调用未定义的num_range也会触发报错

可运行代码 #

# -*- coding: utf-8 -*-
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

# 读取数据保留指定列
df = pd.read_csv('DATA.csv', usecols=['EID', 'ROTATION1','ROTATION2', 'ROTATION3', 'ROTATION4', 'ROTATION5'])

# 构建值到数值的映射字典，按要求的顺序排列
cat_order = ['ROT', 'STILL', 'FLIP'] + list(range(160, 181, 2))
val_map = {val: idx for idx, val in enumerate(cat_order)}
# 兼容数据里数值存为字符串的情况
for num in range(160, 181, 2):
    val_map[str(num)] = val_map[num]

# 将数据转为映射后的数值，空值保留
rot_cols = ['ROTATION1','ROTATION2', 'ROTATION3', 'ROTATION4', 'ROTATION5']
plot_data = df[rot_cols].applymap(lambda x: val_map.get(x, np.nan))
# 设置EID为行索引
plot_data.index = df['EID']

# 绘制热力图，可根据显示需求调整figsize参数
plt.figure(figsize=(6, 15))
# bad参数设置空值位置为透明无填充，yticklabels设为10即每10行显示一个EID，可改为20
ax = sns.heatmap(plot_data, cmap='viridis', cbar=True, yticklabels=10, mask=plot_data.isnull())

# 调整颜色条标签为原始分类值，顺序符合要求
cbar = ax.collections[0].colorbar
cbar.set_ticks(range(len(cat_order)))
cbar.set_ticklabels(cat_order)

# 调整轴显示样式
ax.set_xlabel('')
ax.set_ylabel('EID')
plt.xticks(rotation=0)
plt.tight_layout()
plt.show()

关键实现说明 #

• 数值映射：按照你要求的图例顺序构建映射字典，把字符串分类和数值统一转为连续整数，确保颜色条顺序符合预期
• 空值处理：未匹配到的值和原始空值统一设为np.nan，配合mask参数会自动保留该行，对应位置无填充色
• y轴刻度：yticklabels参数设为10即每10行显示一个EID，改为20就每20行显示一个，适配200行的大数据集显示需求
• x轴默认就是ROTATION1到ROTATION5共5列，完全符合需求

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