如何基于含trajectory、segment的DataFrame绘制目标频谱图？

一、生成目标DataFrame（df_new） #

先按trajectory和segment分组，取每组prediction的众数，得到去重后的目标数据集：

import pandas as pd

# 原始数据
data = {'trajectory': [101,102,102,102,102,102,102,102,104,104,104,104,104,104,104,107,107,107,107,
          107,107,107,107,107,108,108,108,108,108,108,108,109,109,109,109,109,109,112,
         112,112,112,112,113,113,113,113,114,114,114,114],
 'segment': [1,1,1,1,2,2,3,3,1,1,2,2,2,3,3,1,1,2,2,2,2,3,3,3,1,1,1,
          2,2,2,2,1,1,1,2,2,2,1,1,2,2,2,1,2,2,3,1,2,2,2],
  'prediction': [3,0,0,1,3,3,2,2,0,0,4,4,2,0,0,0,0,2,2,2,3,0,0,2,0,0,1,1,
          1,1,0,1,2,1,3,3,3,1,1,4,4,2,1,4,4,3,0,3,3,2]}

df = pd.DataFrame(data)

# 分组计算prediction的众数，取第一个众数结果
df_new = df.groupby(['trajectory', 'segment'])['prediction'].agg(lambda x: x.mode().iloc[0]).reset_index()
print(df_new)

运行后输出的df_new与示例完全一致。

二、绘制指定样式的频谱图 #

使用seaborn热图功能实现，步骤如下：

1. 转换数据为透视表格式 #

将长格式的df_new转为宽格式，适配热图绘制需求：

pivot_df = df_new.pivot(index='trajectory', columns='segment', values='prediction')

2. 定义预测值与颜色的映射 #

为每个prediction标签分配对应颜色：

color_map = {0: '#1f77b4', 1: '#ff7f0e', 2: '#2ca02c', 3: '#d62728', 4: '#9467bd'}
# 按预测值排序，确保颜色顺序对应
sorted_colors = [color_map[i] for i in sorted(color_map.keys())]

3. 绘制频谱图 #

import seaborn as sns
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 8))
# 绘制热图，添加数值标注，设置颜色条
sns.heatmap(pivot_df, cmap=sorted_colors, 
            annot=True, fmt='d', cbar=True,
            xticklabels=pivot_df.columns, yticklabels=pivot_df.index,
            cbar_kws={'ticks': sorted(color_map.keys())})

# 设置坐标轴与标题
plt.xlabel('Segment')
plt.ylabel('Trajectory')
plt.title('Trajectory-Segment Prediction Spectrum')

plt.show()

最终生成的图中：

• Y轴为trajectory，X轴为segment
• 每个单元格颜色对应该轨迹-分段组合的预测众数
• 单元格内标注具体的预测值
• 右侧颜色条展示预测值与颜色的对应关系

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