python - 从交互式 matplotlib 图中获取数据

Question

我正在尝试使用带有 matplotlib 的交互式绘图来手动更正自动生成的散点图数据。（该程序试图检测时间序列数据中的峰值并在每个峰值上放置一个点，但时间序列数据中存在一些不恰当地检测到峰值的伪影。）

我已经成功编写了绘制时间序列并覆盖散点图数据的代码，然后用户可以通过单击添加或删除点。但是，我希望脚本将编辑的点系列保存为新数组并将其传递给其他函数等。没有鼠标单击事件我无法真正调用该函数，所以我不知道如何返回数据。

代码改编自这里：Interactively add and remove scatter points in matplotlib

import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(figsize=(16,4))
a_input = np.sin(range(100))*np.random.normal(20,10,100)
b_input = [ 5, 15, 25, 30, 40, 50, 75, 85]

a = plt.plot(range(len(a_input)),a_input,color='red')[0]
b = plt.scatter(b_input,a_input[b_input],color='grey',s=50,picker=5)
newpeaks = np.array([])

def add_or_remove_point(event):
    global a
    xydata_a = np.stack(a.get_data(),axis=1)
    xdata_a = a.get_xdata()
    ydata_a = a.get_ydata()
    global b
    xydata_b = b.get_offsets()
    xdata_b = b.get_offsets()[:,0]
    ydata_b = b.get_offsets()[:,1]
    global newpeaks

    #click x-value
    xdata_click = event.xdata
    #index of nearest x-value in a
    xdata_nearest_index_a = (np.abs(xdata_a-xdata_click)).argmin()
    #new scatter point x-value
    new_xdata_point_b = xdata_a[xdata_nearest_index_a]
    #new scatter point [x-value, y-value]
    new_xydata_point_b = xydata_a[new_xdata_point_b,:]

    if event.button == 1:
        if new_xdata_point_b not in xdata_b:

            #insert new scatter point into b
            new_xydata_b = np.insert(xydata_b,0,new_xydata_point_b,axis=0)
            #sort b based on x-axis values
            new_xydata_b = new_xydata_b[np.argsort(new_xydata_b[:,0])]
            #update b
            b.set_offsets(new_xydata_b)
            newpeaks = b.get_offsets()[:,0]
            plt.draw()
    elif event.button == 3:
        if new_xdata_point_b in xdata_b:
            #remove xdata point b
            new_xydata_b = np.delete(xydata_b,np.where(xdata_b==new_xdata_point_b),axis=0)
            #print(new_xdata_point_b)
            #update b
            b.set_offsets(new_xydata_b)
            newpeaks = b.get_offsets()[:,0]
        plt.draw()

    print(len(newpeaks))
    return newpeaks

fig.canvas.mpl_connect('button_press_event',add_or_remove_point)
plt.show()

有趣的是，b.get_offsets()确实在函数之外保持更新，所以如果我在脚本完成后在命令行中手动重新分配 newpeaks，我会得到正确的更新数据。但是如果我把作业

newpeaks = b.get_offsets()[:,0]

在脚本中的函数之后，它首先运行并且在绘图关闭时不会更新。我希望它在绘图关闭后立即自动更新，或者以某种方式更新全局变量，而不必手动进行。欢迎任何建议！

标签： pythonmatplotlib