如何在JupyterLab中提取plotly.express图表选中数据的索引或掩码

直接调用fig.show()渲染的是静态HTML输出，图表的交互选中状态不会同步回Python内核，因此读取fig.data[0].selectedpoints永远无法拿到最新的选中数据标识，需要使用Plotly提供的FigureWidget组件实现状态双向同步。

import plotly.express as px
import plotly.graph_objects as go
from ipywidgets import Output
from IPython.display import display

# 加载测试数据集
df = px.data.iris()

# 创建FigureWidget，传入原始索引作为自定义数据，方便后续映射回原DataFrame
fig = px.scatter(df, x="sepal_width", y="sepal_length", color="species", custom_data=[df.index])
fw = go.FigureWidget(fig)

# 定义选中事件回调
out = Output()
selected_mask = [False] * len(df) # 全局布尔掩码，可直接用于筛选原DataFrame
selected_df = None # 全局选中子集变量

@out.capture(clear_output=True)
def on_selection(trace, points, selector):
    global selected_mask, selected_df
    # 重置掩码
    selected_mask = [False] * len(df)
    # 遍历所有选中点，拿到原始索引更新掩码
    for idx in points.customdata:
        selected_mask[idx[0]] = True
    # 筛选得到选中子集
    selected_df = df[selected_mask]
    # 打印选中数据信息，可按需删除
    print(f"共选中{len(selected_df)}条数据")
    display(selected_df)

# 为所有trace绑定选中事件
for trace in fw.data:
    trace.on_selection(on_selection)

# 渲染图表和输出组件
display(fw, out)

• 运行代码后会渲染出和原示例一致的交互式散点图
• 使用图表工具栏的框选/套索工具选中目标区域，选中完成后会自动输出对应的数据子集
• 全局变量selected_mask就是你需要的布尔掩码，直接对原DataFrame切片即可得到选中数据
• 全局变量selected_df已经是筛选后的数据集，可直接用于后续分析

• 首次运行前需要确保JupyterLab已安装ipywidgets依赖，执行以下命令安装即可：

  pip install ipywidgets==7.7.5 plotly==4.14.3
  jupyter labextension install @jupyter-widgets/jupyterlab-manager plotlywidget@4.14.3
  

• 如果不需要实时触发回调，也可以手动读取选中状态：框选完成后直接执行fw.data[0].selectedpoints即可拿到当前trace的选中点下标，配合customdata映射到原始索引即可。

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