Matplotlib中如何保留折线图配色并自定义含分组信息的图例？

问题描述 #

我在同一坐标轴上绘制了两组折线图，每组对应一个分组，包含代表两种条件的两条折线：Group 1使用Blues配色方案，Group 2使用Reds配色方案。自动生成的图例仅保留了各折线的配色，但缺失分组信息。

当尝试通过以下代码手动修改图例以添加分组归属时：

plt.legend(title='Legend', labels=['Group 1: C1', 'Group 1: C2', 'Group 2: C1', 'Group 2: C2'])

图例中的其中一组配色信息丢失。现寻求解决方法：如何在手动修改图例时保留两组折线的原有配色，或是否有其他方法在Matplotlib中为图例添加分组维度？

所用数据集：

data = {'group': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
        'block': [1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3],
        'cond': ['c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2', 'c1', 'c2'],
        'value': [1, 2, 1.2, 2.1, 1.1, 2.2, 3, 4, 3.3, 4.1, 3.1, 4]}

原绘图代码：

sns.lineplot(data=df.loc[df['group'] == 1], x='block', y='value', hue='cond',
             palette='Blues', marker='o')

sns.lineplot(data=df.loc[df['group'] == 2], x='block', y='value', hue='cond',
             palette='Reds', marker='o')
plt.show()

解决方法 #

方法1：手动修正图例，保留原有配色 #

直接调用plt.legend()会重置整个图例，覆盖Seaborn生成的原有图例元素。正确做法是先获取当前图例的所有图形元素（handles），再修改标签：

import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

df = pd.DataFrame(data)

# 绘制两组折线
sns.lineplot(data=df.loc[df['group'] == 1], x='block', y='value', hue='cond',
             palette='Blues', marker='o')
sns.lineplot(data=df.loc[df['group'] == 2], x='block', y='value', hue='cond',
             palette='Reds', marker='o')

# 获取当前图例的handles（图形元素）
handles, _ = plt.gca().get_legend_handles_labels()
# 设置新的标签文本
new_labels = ['Group 1: C1', 'Group 1: C2', 'Group 2: C1', 'Group 2: C2']
# 更新图例，传入原handles和新labels
plt.legend(title='Legend', handles=handles, labels=new_labels)

plt.show()

方法2：通过组合列自动生成带分组的图例 #

先给数据集添加一个分组+条件的组合列，将分组和条件信息合并，再用这个列作为hue参数，并指定对应配色，Seaborn会自动生成包含分组信息的图例：

import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

df = pd.DataFrame(data)
# 创建分组与条件的组合列，格式化显示文本
df['group_cond'] = df.apply(lambda row: f"Group {row['group']}: {row['cond'].upper()}", axis=1)

# 自定义配色：Group1的C1/C2取Blues调色板的颜色，Group2取Reds调色板
palette = {
    'Group 1: C1': sns.color_palette('Blues')[2],
    'Group 1: C2': sns.color_palette('Blues')[4],
    'Group 2: C1': sns.color_palette('Reds')[2],
    'Group 2: C2': sns.color_palette('Reds')[4]
}

# 一次性绘制所有折线，用组合列作为hue
sns.lineplot(data=df, x='block', y='value', hue='group_cond',
             palette=palette, marker='o')

plt.legend(title='Legend')
plt.show()

方法3：构建嵌套分组图例 #

如果想要更清晰的分组展示（比如先按Group分类，再按Cond子分类），可以手动创建嵌套的Matplotlib图例：

import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.legend import Legend
import pandas as pd

df = pd.DataFrame(data)

# 分别绘制四条折线，并保存每个折线的对象
line1_c1, = sns.lineplot(data=df.loc[(df['group'] ==1) & (df['cond']=='c1')], 
                         x='block', y='value', color=sns.color_palette('Blues')[2], marker='o')
line1_c2, = sns.lineplot(data=df.loc[(df['group'] ==1) & (df['cond']=='c2')], 
                         x='block', y='value', color=sns.color_palette('Blues')[4], marker='o')
line2_c1, = sns.lineplot(data=df.loc[(df['group'] ==2) & (df['cond']=='c1')], 
                         x='block', y='value', color=sns.color_palette('Reds')[2], marker='o')
line2_c2, = sns.lineplot(data=df.loc[(df['group'] ==2) & (df['cond']=='c2')], 
                         x='block', y='value', color=sns.color_palette('Reds')[4], marker='o')

# 获取当前画布和坐标轴对象
fig, ax = plt.gcf(), plt.gca()

# 创建Group1的子图例
leg1 = Legend(ax, [line1_c1, line1_c2], ['C1', 'C2'], title='Group 1', 
              loc='upper left', bbox_to_anchor=(1, 1))
# 创建Group2的子图例
leg2 = Legend(ax, [line2_c1, line2_c2], ['C1', 'C2'], title='Group 2', 
              loc='upper left', bbox_to_anchor=(1, 0.7))

# 将子图例添加到坐标轴
ax.add_artist(leg1)
ax.add_artist(leg2)

plt.show()

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