告别Matplotlib!Python量化分析神器PyQtGraph入门指南

import sysimport pandas as pdfrom PyQt6.QtWidgets import QApplication, QLabel, QMainWindow, QMenu, QStatusBar, QToolBar, QVBoxLayout, QWidgetfrom PyQt6.QtCore import Qtfrom PyQt6.QtGui import QAction, QKeySequenceimport pyqtgraph as pgimport numpy as np

class CandlestickItem(pg.GraphicsObject):    def __init__(self, data):        pg.GraphicsObject.__init__(self)        self.data = data        self.generatePicture()    def generatePicture(self):        """生成K线图绘制指令，优化性能的关键"""        self.picture = pg.QtGui.QPicture()        p = pg.QtGui.QPainter(self.picture)        p.setPen(pg.mkPen("w"))        # 计算K线宽度        w = 0.6 if len(self.data) > 1 else 0.8        for (t, open_price, close_price, low_price, high_price) in self.data:            # 涨跌颜色：红色表示跌，绿色表示涨            if open_price > close_price:                p.setBrush(pg.mkBrush("r"))            else:                p.setBrush(pg.mkBrush("g"))            # 绘制K线实体            rect_top = min(open_price, close_price)            rect_height = abs(close_price - open_price)            p.drawRect(pg.QtCore.QRectF(                t - w/2,     # 矩形左边界                rect_top,    # 矩形顶部                w,           # 矩形宽度                rect_height  # 矩形高度            ))            # 绘制上下影线            p.drawLine(                pg.QtCore.QPointF(t, low_price),                pg.QtCore.QPointF(t, high_price)            )        p.end()    def paint(self, p, *args):        p.drawPicture(0, 0, self.picture)    def boundingRect(self):        return pg.QtCore.QRectF(self.picture.boundingRect())

class Crosshair:    def __init__(self, viewer, kline_plot, volume_plot):        self.viewer = viewer        self.kline_plot = kline_plot        self.volume_plot = volume_plot        # 创建十字线        self.vLine = pg.InfiniteLine(            angle=90, movable=False,             pen=pg.mkPen("yellow", width=1)        )        self.hLine = pg.InfiniteLine(            angle=0, movable=False,            pen=pg.mkPen("yellow", width=1)        )        # 添加到图表        kline_plot.addItem(self.vLine, ignoreBounds=True)        kline_plot.addItem(self.hLine, ignoreBounds=True)        volume_plot.addItem(self.vLine, ignoreBounds=True)        # 信息标签        self.label = pg.TextItem(anchor=(0, 1))        self.label.setParentItem(kline_plot.getViewBox())        self.label.setVisible(False)        kline_plot.addItem(self.label)        # 连接鼠标事件        self.kline_plot.scene().sigMouseMoved.connect(self.mouse_moved)        self.kline_plot.scene().sigMouseClicked.connect(self.mouse_clicked)    def mouse_moved(self, pos):        """实时更新十字光标位置和数据显示"""        if not self.enabled:            return        mouse_point = self.kline_plot.getViewBox().mapSceneToView(pos)        x_range = self.kline_plot.getViewBox().viewRange()[0]        y_range = self.kline_plot.getViewBox().viewRange()[1]        # 检查鼠标是否在视图内        if x_range[0] <= mouse_point.x() <= x_range[1] and y_range[0] <= mouse_point.y() <= y_range[1]:            self.vLine.setPos(mouse_point.x())            self.hLine.setPos(mouse_point.y())            # 显示详细信息            x_val = int(round(mouse_point.x()))            y_val = float(f"{mouse_point.y():.4f}")            if 0 <= x_val < len(self.viewer.date_labels):                date = self.viewer.date_labels.iloc[x_val]                date_str = date.strftime("%Y-%m-%d") if hasattr(date, "strftime") else str(date)                self.label.setText(f"Date:{date_str}\nPrice:{y_val}", color="w")                self.label.setPos(mouse_point.x(), mouse_point.y())                self.label.setVisible(True)    def mouse_clicked(self, evt):        """右键点击切换光标显示"""        if evt.button() == Qt.MouseButton.RightButton:            self.enabled = not self.enabled            self.vLine.setVisible(self.enabled)            self.hLine.setVisible(self.enabled)            self.label.setVisible(self.enabled)

class StockChartViewer(QMainWindow):    def __init__(self, data):        super().__init__()        self.data = data.copy()        self.date_labels = []        self.crosshair = None        self.init_ui()        self.init_toolbar()    def init_ui(self):        """初始化用户界面"""        self.setWindowTitle("股票数据分析")        self.setGeometry(100, 100, 1400, 900)        # 设置中心部件        central_widget = QWidget()        self.setCentralWidget(central_widget)        main_layout = QVBoxLayout(central_widget)        main_layout.setContentsMargins(0, 0, 0, 0)        # 创建图形布局部件        self.graphics_widget = pg.GraphicsLayoutWidget()        main_layout.addWidget(self.graphics_widget)        # 状态栏        self.status_bar = QStatusBar()        self.setStatusBar(self.status_bar)        self.status_label = QLabel(            "就绪 | 左键：拖动 | 右键：菜单 | 滚轮：缩放 | 右键双击：重置 | 右键点击：切换十字光标"        )        self.status_bar.addWidget(self.status_label)        # 启用抗锯齿        pg.setConfigOptions(antialias=True)        self.plot_charts()    def plot_charts(self):        """绘制K线图和交易量图"""        # 处理日期数据        if "date" in self.data.columns:            if not pd.api.types.is_datetime64_any_dtype(self.data["date"]):                self.data["date"] = pd.to_datetime(self.data["date"])            self.date_labels = self.data["date"].copy()            self.dates_numeric = np.arange(len(self.data))        # 获取数据        dates_numeric = self.dates_numeric        o = self.data["open"].values        h = self.data["high"].values        l = self.data["low"].values        c = self.data["close"].values        volume = self.data["amount"].values        # 创建K线图        self.kline_plot = self.graphics_widget.addPlot(            title="K线图", row=0, col=0        )        self.kline_plot.addLegend()        # 设置坐标轴范围        if len(dates_numeric) > 0:            x_min, x_max = dates_numeric.min(), dates_numeric.max()            self.kline_plot.setXRange(x_min, x_max, padding=0.1)            price_min = min(o.min(), h.min(), l.min(), c.min())            price_max = max(o.max(), h.max(), l.max(), c.max())            self.kline_plot.setYRange(price_min * 0.95, price_max * 1.05, padding=0.1)        # 创建K线数据        candlestick_data = []        for i in range(len(dates_numeric)):            candlestick_data.append((                float(dates_numeric[i]),                 float(o[i]),                 float(c[i]),                 float(l[i]),                 float(h[i])            ))        candlestick_item = CandlestickItem(candlestick_data)        self.kline_plot.addItem(candlestick_item)        # 设置图表属性        self.kline_plot.setLabel("left", "价格")        self.kline_plot.setLabel("bottom", "日期")        self.kline_plot.showGrid(x=True, y=True, alpha=0.3)        self.kline_plot.setMouseEnabled(x=True, y=True)        # 设置X轴标签        self.set_x_axis_labels(self.kline_plot)        # 创建交易量图        self.volume_plot = self.add_volume_chart(            dates_numeric, volume, o, c, 1, self.kline_plot        )        # 初始化十字光标        self.crosshair = Crosshair(self, self.kline_plot, self.volume_plot)        self.kline_plot.scene().sigMouseMoved.connect(self.update_status_bar)    def add_volume_chart(self, dates, volume, open_prices, close_prices, row, linked_plot):        """添加交易量图"""        volume_plot = self.graphics_widget.addPlot(title="交易量", row=row, col=0)        volume_plot.setXLink(linked_plot)  # 链接X轴        volume_plot.setMouseEnabled(x=True, y=True)        # 绘制交易量柱状图        bar_width = 0.6 if len(dates) > 1 else 0.8        for i, (date, vol, open_price, close_price) in enumerate(            zip(dates, volume, open_prices, close_prices)        ):            color = "g" if close_price >= open_price else "r"            brush = pg.mkBrush(color)            bar = pg.BarGraphItem(                x=[date], height=[vol], width=bar_width, brush=brush            )            volume_plot.addItem(bar)        # 设置坐标轴        self.set_x_axis_labels(volume_plot)        volume_plot.setLabel("left", "交易量")        volume_plot.setLabel("bottom", "日期")        volume_plot.showGrid(x=True, y=True, alpha=0.3)        return volume_plot

def init_toolbar(self):    """初始化工具栏，提供丰富的交互功能"""    toolbar = QToolBar("主工具栏")    self.addToolBar(toolbar)    # 重置视图按钮    reset_action = QAction("重置视图", self)    reset_action.setStatusTip("重置视图到初始状态")    reset_action.triggered.connect(self.reset_view)    reset_action.setShortcut("R")    toolbar.addAction(reset_action)    toolbar.addSeparator()    # 十字光标开关    crosshair_action = QAction("十字光标", self)    crosshair_action.setStatusTip("显示/隐藏十字光标")    crosshair_action.setCheckable(True)    crosshair_action.setChecked(True)    crosshair_action.triggered.connect(self.toggle_crosshair)    crosshair_action.setShortcut("C")    toolbar.addAction(crosshair_action)    toolbar.addSeparator()    # 缩放模式    zoom_x_action = QAction("X轴缩放", self)    zoom_x_action.setStatusTip("启用X轴缩放")    zoom_x_action.setCheckable(True)    zoom_x_action.setChecked(False)    zoom_x_action.triggered.connect(lambda: self.set_zoom_mode("x"))    toolbar.addAction(zoom_x_action)    zoom_y_action = QAction("Y轴缩放", self)    zoom_y_action.setStatusTip("启用Y轴缩放")    zoom_y_action.setCheckable(True)    zoom_y_action.setChecked(False)    zoom_y_action.triggered.connect(lambda: self.set_zoom_mode("y"))    toolbar.addAction(zoom_y_action)    zoom_xy_action = QAction("XY轴缩放", self)    zoom_xy_action.setStatusTip("启用XY轴缩放")    zoom_xy_action.setCheckable(True)    zoom_xy_action.setChecked(True)    zoom_xy_action.triggered.connect(lambda: self.set_zoom_mode("xy"))    toolbar.addAction(zoom_xy_action)    toolbar.addSeparator()    # 鼠标模式    pan_action = QAction("平移模式", self)    pan_action.setStatusTip("启用平移模式")    pan_action.setCheckable(True)    pan_action.setChecked(True)    pan_action.triggered.connect(lambda: self.set_mouse_mode("pan"))    pan_action.setShortcut("P")    toolbar.addAction(pan_action)    zoom_action = QAction("缩放模式", self)    zoom_action.setStatusTip("启用缩放模式")    zoom_action.setCheckable(True)    zoom_action.setChecked(False)    zoom_action.triggered.connect(lambda: self.set_mouse_mode("zoom"))    zoom_action.setShortcut("Z")    toolbar.addAction(zoom_action)    toolbar.addSeparator()    # 显示设置    show_grid_action = QAction("显示网格", self)    show_grid_action.setStatusTip("显示/隐藏网格")    show_grid_action.setCheckable(True)    show_grid_action.setChecked(True)    show_grid_action.triggered.connect(self.toggle_grid)    toolbar.addAction(show_grid_action)

def create_context_menu(self, ev):    """创建右键上下文菜单"""    menu = QMenu()    reset_action = menu.addAction("重置视图（R）")    reset_action.triggered.connect(self.reset_view)    menu.addSeparator()    # 十字光标选项    crosshair_action = menu.addAction("切换十字光标 (C)")    crosshair_action.setCheckable(True)    crosshair_action.setChecked(self.crosshair.enabled)    crosshair_action.triggered.connect(self.toggle_crosshair)    menu.addSeparator()    # 缩放选项    menu.addAction("X轴缩放").triggered.connect(lambda: self.set_zoom_mode('x'))    menu.addAction("Y轴缩放").triggered.connect(lambda: self.set_zoom_mode('y'))    menu.addAction("XY轴缩放").triggered.connect(lambda: self.set_zoom_mode('xy'))    menu.addSeparator()    # 鼠标模式    menu.addAction("平移模式 (P)").triggered.connect(lambda: self.set_mouse_mode('pan'))    menu.addAction("缩放模式 (Z)").triggered.connect(lambda: self.set_mouse_mode('zoom'))    menu.addSeparator()    # 网格显示    grid_action = menu.addAction("显示网格")    grid_action.setCheckable(True)    grid_action.setChecked(self.kline_plot.ctrl.gridCheck.isChecked())    grid_action.triggered.connect(self.toggle_grid)    return menudef keyPressEvent(self, event):    """键盘快捷键支持"""    if event.key() == Qt.Key.Key_R:        self.reset_view()    elif event.key() == Qt.Key.Key_C:        self.toggle_crosshair()    elif event.key() == Qt.Key.Key_Z:        self.set_mouse_mode('zoom')    elif event.key() == Qt.Key.Key_P:        self.set_mouse_mode('pan')    elif event.key() == Qt.Key.Key_G:        self.toggle_grid()    elif event.key() == Qt.Key.Key_Escape:        self.reset_view()    else:        super().keyPressEvent(event)

def main():    app = QApplication(sys.argv)    app.setStyle("Fusion")  # 使用Fusion风格，跨平台一致性更好    # 假设从get_stock_data模块导入数据    from get_stock_data import df    viewer = StockChartViewer(df)    viewer.show()    sys.exit(app.exec())if __name__ == "__main__":    main()