matplotlib画一条可拖动的曲线

可以修改draw_line传入的n值，决定画几条曲线

import matplotlib
#此处在打包时会用到
matplotlib.use('TkAgg')
import sys,os
import time
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np
from scipy.interpolate import interp1d
from matplotlib.widgets import Button


# 用于创建可拖动的线条对象，并处理鼠标事件。
class DraggableLine:
    def __init__(self, ax, x, y):
        self.ax = ax
        self.x = x
        self.y = y
        self.line, = ax.plot(x, y, marker='o', markersize=5, markerfacecolor='red', linestyle='-', picker=5)
        self.cid_click = self.line.figure.canvas.mpl_connect('button_press_event', self.on_click)
        self.cid_release = self.line.figure.canvas.mpl_connect('button_release_event', self.on_release)
        self.cid_motion = self.line.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)
        self.cid_dblclick = self.line.figure.canvas.mpl_connect('button_press_event', self.on_dblclick)
        self.selected_point = None
        self.annotation = self.ax.annotate('', xy=(0, 0), xytext=(10, 10),
                                           textcoords='offset points', bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.5),
                                           arrowprops=dict(arrowstyle='->'))
        self.annotation.set_visible(False)
        self.selected_points = []
        self.newx = []
        self.newy = []
        self.modify_xdata = []
        self.modify_ydata = []

    # 检测鼠标点击事件，判断是否选中某个点
    def on_click(self, event):
        # 判断事件是否在自己ax区间，不是发生在自己ax区间，不理会
        if event.inaxes != self.ax:
            return
        if event.inaxes != self.line.axes: return
        contains, attr = self.line.contains(event)
        if contains:
            self.selected_point = attr['ind'][0]
        else:
            self.selected_point = None
        self.update_annotation(event)

    # 释放鼠标按钮时，取消选中点。
    def on_release(self, event):
        # 判断事件是否在自己ax区间，不是发生在自己ax区间，不理会
        if event.inaxes != self.ax:
            return
        self.selected_point = None
        xdata = list(self.line.get_xdata())
        ydata = list(self.line.get_ydata())
        self.line.set_data(xdata, ydata)
        self.line.figure.canvas.draw()
        self.update_curve()
        self.annotation.set_visible(False)

    # 拖动鼠标时，更新选中点的纵坐标。
    def on_motion(self, event):
        # 判断事件是否在自己ax区间，不是发生在自己ax区间，不理会
        if event.inaxes != self.ax:
            return
        if self.selected_point is None: return
        if event.inaxes != self.line.axes: return
        ydata = self.line.get_ydata()
        ydata[self.selected_point] = int(event.ydata)
        if self.selected_point not in self.selected_points:
            self.selected_points.append(self.selected_point)
        self.line.set_ydata(ydata)
        self.line.figure.canvas.draw()
        self.update_annotation(event)
        self.update_curve()

    # 更新鼠标悬停时显示的坐标注释。
    def update_annotation(self, event):
        if event.inaxes != self.ax:
            return
        x = self.line.get_xdata()[self.selected_point]
        y = self.line.get_ydata()[self.selected_point]
        # self.annotation.xy = (event.xdata, event.ydata)
        self.annotation.xy = (x, y)
        text = f'({x}, {y})'
        self.annotation.set_text(text)
        self.annotation.set_visible(True)
        self.line.figure.canvas.draw()

    # 检测鼠标双击事件，左键双击添加锚点，右键双击移除锚点。
    def on_dblclick(self, event):
        # 判断事件是否在自己ax区间，不是发生在自己ax区间，不理会
        if event.inaxes != self.ax:
            return
        if event.dblclick:
            if event.button == 1:  # Left double click to add point
                self.add_point(int(event.xdata), int(event.ydata))
            elif event.button == 3:  # Right double click to remove point
                self.remove_point(event.xdata)
            xdata = list(self.line.get_xdata())
            ydata = list(self.line.get_ydata())
            self.line.set_data(xdata, ydata)
            self.line.figure.canvas.draw()
            self.update_curve()

    # 添加新的锚点，并保持锚点按横坐标排序。
    def add_point(self, x, y):
        xdata = list(self.line.get_xdata())
        ydata = list(self.line.get_ydata())
        xdata.append(x)
        ydata.append(y)
        sorted_indices = np.argsort(xdata)
        xdata = np.array(xdata)[sorted_indices]
        ydata = np.array(ydata)[sorted_indices]
        self.line.set_data(xdata, ydata)
        self.update_curve()

    # 移除指定横坐标的锚点
    def remove_point(self, x):
        xdata, ydata = self.get_data()
        distances = np.abs(np.array(xdata) - x)
        # 获取最小值对应的索引
        closest_index = np.argmin(distances)
        if distances[closest_index] < 0.5:  # 如果距离足够近
            xdata.pop(closest_index)
            ydata.pop(closest_index)
            self.line.set_data(xdata, ydata)
            self.update_curve()

    # 更新插值曲线，使用三次样条插值（cubic interpolation）连接锚点。
    def update_curve(self):
        xdata, ydata = self.get_data()
        if len(xdata) > 1:
            f = interp1d(xdata, ydata, kind=3)
            xnew = np.linspace(min(xdata), max(xdata), num=257, endpoint=True)
            ynew = f(xnew)
            self.modify_xdata = xnew
            self.modify_ydata = ynew
            self.ax.lines[-1].set_data(xnew, ynew)
            self.line.figure.canvas.draw()

    # 获取当前线条的坐标数据
    def get_data(self):
        return self.line.get_xdata(), self.line.get_ydata()

# xx 是横坐标数组
# yy是纵坐标数组
# x_points 是x轴锚点数组
# y_points 是y轴锚点数组
def draw_line(x, yy, x_point, y_points, n):
    # 设置绘图面板的显示尺寸9*100
    plt.rcParams['figure.figsize'] = (6, 6)
    # 纵向排列
    fig, ax = plt.subplots(n, 1)
    print(n)
    if n == 1:
        ax.set_title(rf'Draggable Line {0}')
        ax.set_xlim(-5, 12)
        ax.set_ylim(-5, 100)
    else:
        for i in range(0, n):
            ax[i].set_title(rf'Draggable Line {i}')
            ax[i].set_xlim(-5, 12)
            ax[i].set_ylim(-5, 100)

    if 1 == n:
        original_line = ax.plot(x, yy[0], linestyle='--', color='blue', label=f'Original Line{0}')
        draggable_lines = DraggableLine(ax, x_point, y_points[0])
        p_line, = ax.plot(x_point, y_points[0], linestyle='-', color='red', label=f'Modified line{0}')
        plt.legend()
        plt.show()
    else:
        # 画原始曲线（虚线）
        original_lines = [0] * n
        for i in range(0, n):
            original_lines[i], = ax[i].plot(x, yy[i], linestyle='--', color='blue', label=f'Original line{i}')

        # 创建可拖动的线条
        draggable_lines = [0] * n
        red_lines = [0] * n
        for i in range(0, n):
            # 添加一条用于显示插值曲线的实线
            draggable_lines[i] = DraggableLine(ax[i], x_point, y_points[i])
            p_line, = ax[i].plot(x_point, y_points[i], linestyle='-', color='red', label=f'Modified line{i}')

        plt.legend()
        plt.show()
if __name__ == '__main__':
    xx = [0,1,2,3,4,5,6,7,8,9]
    yy = []
    y1 = [-1,2,5,8,11,14,17,20,23,26]
    yy.append(y1)
    y2 = [20, 30, 35, 40, 45, 50, 55, 65, 75, 80]
    yy.append(y2)
    x_points = [0,2,4,6,8]
    y_points = []
    y_point1 = [-1,5,11,17,23]
    y_points.append(y_point1)
    y_point2 = [20,35,45,55,75]
    y_points.append(y_point2)
    draw_line(xx, yy, x_points, y_points, 2)