如何在matplotlib图例中添加可显示颜色渐变效果的LineCollection项

直接重写LineCollection对应的图例处理器，在图例区域生成一小段匹配原线条配色规则的渐变LineCollection作为标识，替换默认处理器生成的单色Line2D即可，不需要改动原有渐变线的绘制逻辑。

核心实现代码 #

自定义图例处理器：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.legend_handler import HandlerLineCollection

class GradientLineLegendHandler(HandlerLineCollection):
    def create_artists(self, legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans):
        # 图例区域内生成水平居中的线段采样点
        x_samples = np.linspace(xdescent, xdescent + width, 100)
        y_samples = np.full_like(x_samples, height/2 - ydescent)
        line_points = np.column_stack([x_samples, y_samples])
        # 组装为LineCollection要求的段结构
        segments = np.array([line_points[i:i+2] for i in range(len(line_points)-1)])
        # 复用原线条的配色映射、线宽、透明度等属性
        legend_line = LineCollection(
            segments,
            cmap=orig_handle.get_cmap(),
            norm=orig_handle.norm,
            linewidth=orig_handle.get_linewidth(),
            alpha=orig_handle.get_alpha()
        )
        # 配置渐变范围，默认展示完整色阶
        legend_line.set_array(np.linspace(0, 1, len(segments)))
        legend_line.set_transform(trans)
        return [legend_line]

使用方式 #

给legend方法传入handler映射，将LineCollection类型绑定到上面自定义的处理器即可：

# 以下为测试代码
if __name__ == "__main__":
    fig, ax = plt.subplots(figsize=(7, 4))
    # 绘制测试用渐变正弦曲线
    t = np.linspace(0, 4*np.pi, 200)
    x, y = t, np.sin(t)
    points = np.column_stack([x, y])
    segments = np.array([points[i:i+2] for i in range(len(points)-1)])
    gradient_line = LineCollection(segments, cmap="coolwarm", linewidth=3)
    gradient_line.set_array(t) # 沿x轴映射渐变色
    ax.add_collection(gradient_line)
    ax.autoscale()

    # 添加带渐变效果的图例
    ax.legend(
        [gradient_line],
        ["示例渐变曲线"],
        handler_map={LineCollection: GradientLineLegendHandler()},
        loc="upper right"
    )
    plt.show()

适配调整说明 #

• 如果原线条是沿y值、其他自定义维度做颜色映射，只需要修改set_array传入的数组规则，和原线条的配色映射逻辑保持一致即可。
• 如果只需要展示部分色段，调整np.linspace的起止范围即可，比如传入np.linspace(0.2, 0.7, len(segments))就会只展示原色条20%到70%位置的渐变。
• 代码兼容matplotlib 3.4及以上版本，更低版本需要微调create_artists的入参适配。

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