如何为Matplotlib散点图启用Legend picking（图例拾取）功能

Matplotlib官方提供的折线图图例点击切换示例仅支持折线图——leg.get_lines()只能获取折线对应的图例元素，散点图的图例元素无法通过该方法获取。我们需要实现一个仅依赖Axes对象的通用方案，同时支持折线图和散点图的图例点击切换，无需保存绘图时的折线/散点对象引用。

问题分析 #

• 折线图的绘图元素存储在ax.lines中，对应图例元素可通过leg.get_lines()获取，但该方法不包含散点图的图例
• 散点图的绘图元素是matplotlib.collections.PathCollection类型，存储在ax.collections中
• 图例的所有元素（包括折线和散点的代理）可以通过leg.get_handles()统一获取

完整实现代码 #

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections

def setup_legend_picking(ax):
    # 创建图例
    leg = ax.legend(fancybox=True, shadow=True)
    # 获取Axes中所有的折线和散点元素
    plot_elements = []
    # 添加折线
    plot_elements.extend(ax.lines)
    # 添加散点（筛选PathCollection类型）
    plot_elements.extend([elem for elem in ax.collections if isinstance(elem, matplotlib.collections.PathCollection)])
    
    # 建立图例元素与原始绘图元素的映射：通过label匹配
    elem_map = {}
    for handle in leg.get_handles():
        # 找到与图例handle label匹配的原始元素
        for elem in plot_elements:
            if elem.get_label() == handle.get_label():
                elem_map[handle] = elem
                # 开启图例元素的拾取功能
                handle.set_picker(True)
                break
    
    def on_pick(event):
        # 获取被点击的图例元素
        leg_elem = event.artist
        if leg_elem not in elem_map:
            return
        # 获取对应的原始绘图元素
        orig_elem = elem_map[leg_elem]
        # 切换可见性
        visible = not orig_elem.get_visible()
        orig_elem.set_visible(visible)
        # 更新图例元素的透明度，标识状态
        leg_elem.set_alpha(1.0 if visible else 0.2)
        # 重绘画布
        ax.figure.canvas.draw()
    
    # 绑定拾取事件
    ax.figure.canvas.mpl_connect('pick_event', on_pick)
    return leg

# 测试示例
if __name__ == "__main__":
    fig, ax = plt.subplots()
    ax.set_title('点击图例切换对应图形显示/隐藏')
    
    # 添加折线
    t = np.linspace(0, 1)
    ax.plot(t, 2*np.sin(2*np.pi*t), lw=2, label='1 Hz')
    ax.plot(t, 4*np.sin(2*np.pi*2*t), lw=2, label='2 Hz')
    
    # 添加散点
    x = np.random.rand(50)
    y = np.random.rand(50)
    ax.scatter(x, y, label='散点图')
    
    # 设置图例拾取功能
    setup_legend_picking(ax)
    
    plt.show()

关键说明 #

1. 元素收集：从ax.lines获取折线，从ax.collections筛选出散点类型的元素
2. 映射建立：通过元素的label属性匹配图例handle和原始绘图元素，确保一一对应
3. 拾取事件：点击图例元素时，切换对应原始图形的可见性，并通过透明度变化反馈状态
4. 通用性：仅依赖传入的Axes对象，无需保存绘图时的临时变量

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