如何用Matplotlib绘制按group_id变色的单折线图？

示例数据 #

import pandas as pd
df = pd.DataFrame({'x': ['2008-01-01', '2008-01-02', '2008-01-03', '2008-01-04'], 
                   'y': [1, 2, 3, 6],
                   'group_id': ['OBSERVED', 'IMPUTED', 'OBSERVED', 'IMPUTED'], 
                   'color': ['blue', 'red', 'blue', 'red']})
df['x'] = pd.to_datetime(df['x'])

需求 #

• 绘制一条连续折线，而非多条独立线段
• 折线的线段颜色根据group_id（对应color列）自动切换

原代码的问题 #

• 按group_id分组绘图会将同一组的点单独连线，不同组的点之间断开，最终呈现两条独立折线
• 使用data.plot()时，matplotlib会默认使用内置配色方案，忽略color列的自定义颜色

解决方案 #

通过遍历相邻点对，逐段绘制线段并指定对应颜色，即可实现连续且分段变色的折线：

import matplotlib.pyplot as plt
from matplotlib.patches import Patch

fig, ax = plt.subplots(1)

# 逐段绘制相邻点之间的线段
for i in range(len(df) - 1):
    # 获取当前线段的起点和终点数据
    point1 = df.iloc[i]
    point2 = df.iloc[i+1]
    # 用起点的color绘制当前线段
    ax.plot([point1['x'], point2['x']], [point1['y'], point2['y']], color=point1['color'])

# 手动创建图例（匹配group_id与对应颜色）
legend_patches = [
    Patch(facecolor='blue', label='OBSERVED'),
    Patch(facecolor='red', label='IMPUTED')
]
ax.legend(handles=legend_patches)

plt.show()

说明 #

• 循环遍历每一对相邻数据点，单独绘制两点间的线段，确保折线整体连续
• 每段线段的颜色取起点行的color值（因group_id与color一一对应，线段所属组以起点组为准）
• 手动创建图例元素，解决分段绘制无法自动生成组图例的问题

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