Matplotlib 3D散点图pick_event返回错误索引的原因排查

我用Python和Matplotlib绘制3D散点图，想要实现点击某个点时将其高亮显示。参考了Matplotlib官方的简单拾取示例，也查了相关问题，但on_pick函数返回的索引始终对应不上正确的3D散点。

我修改官方示例后的代码如下：

import matplotlib.pyplot as plt
import numpy as np

points_3d = np.random.rand(10, 3)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

scatter = ax.scatter(points_3d[:, 0], points_3d[:, 1], points_3d[:, 2], picker=True, c='blue')

def onpick(event):
    if event.artist != scatter:
        return
    n = len(event.ind)
    if not n:
        return
    for dataind in event.ind:
        pt = points_3d[dataind]
        ax.scatter(pt[0], pt[1], pt[2], c='red', marker='D', s=100)
    fig.canvas.draw()
    return True


fig.canvas.mpl_connect('pick_event', onpick)

plt.show()

实际运行效果：

为什么pick_event返回的索引无法对应到正确的3D散点？

原因分析 #

Matplotlib的3D散点图拾取逻辑和2D存在差异：在3D轴中，散点会根据当前视角进行深度排序（为了正确渲染遮挡关系），导致内部存储的点顺序和你传入的points_3d原始顺序不一致。event.ind返回的是散点绘制顺序的索引，而非原始数据的索引，直接用它去索引原始数组自然会出错。

解决方案 #

方案1：直接使用排序后的坐标绘制高亮

通过散点对象的_offsets3d属性可以获取已经按3D视角排序后的点坐标，用这个坐标绘制高亮就能匹配点击的点：

import matplotlib.pyplot as plt
import numpy as np

points_3d = np.random.rand(10, 3)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

scatter = ax.scatter(points_3d[:, 0], points_3d[:, 1], points_3d[:, 2], picker=True, c='blue')

def onpick(event):
    if event.artist != scatter:
        return
    n = len(event.ind)
    if not n:
        return
    # 获取排序后的散点坐标
    xs, ys, zs = scatter._offsets3d
    for dataind in event.ind:
        ax.scatter(xs[dataind], ys[dataind], zs[dataind], c='red', marker='D', s=100)
    fig.canvas.draw()
    return True


fig.canvas.mpl_connect('pick_event', onpick)

plt.show()

方案2：匹配原始数据索引（如需关联原始数据）

如果需要获取点击点对应的原始数据索引，可以通过坐标匹配的方式处理（注意浮点数精度问题，用np.allclose）：

import matplotlib.pyplot as plt
import numpy as np

points_3d = np.random.rand(10, 3)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

scatter = ax.scatter(points_3d[:, 0], points_3d[:, 1], points_3d[:, 2], picker=True, c='blue')

def onpick(event):
    if event.artist != scatter:
        return
    n = len(event.ind)
    if not n:
        return
    xs, ys, zs = scatter._offsets3d
    for dataind in event.ind:
        picked_pt = np.array([xs[dataind], ys[dataind], zs[dataind]])
        # 匹配原始数据中的对应索引
        original_ind = np.where(np.allclose(points_3d, picked_pt, atol=1e-6))[0][0]
        print(f"点击点对应的原始数据索引：{original_ind}")
        ax.scatter(picked_pt[0], picked_pt[1], picked_pt[2], c='red', marker='D', s=100)
    fig.canvas.draw()
    return True


fig.canvas.mpl_connect('pick_event', onpick)

plt.show()

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