基于pyqtgraph绘图库时频域显示


import numpy as np
import pyqtgraph as pg
from PyQt5.QtCore import Qt
from PyQt5.QtCore import QPointF
from PyQt5.QtCore import QRectF
from PyQt5.QtGui import QColor


chn_color = [
    QColor(Qt.red),
    QColor(Qt.green),
    QColor(70, 130, 180),  # steelBlue
    QColor(Qt.yellow),
    QColor(Qt.cyan),
    QColor(Qt.magenta),
    QColor(255, 140, 0),    # darkOrange
    QColor(255, 182, 193),  # lightPink
    QColor(Qt.blue),
    QColor(255, 105, 180),  # hotPink
    QColor(46, 139, 87),    # seagreen
    QColor(199, 21, 133),   # mediumVioletRed
    QColor(255, 69, 0),		# orangeRed
    QColor(Qt.darkYellow),
    QColor(160, 82, 45),    # sienna
    QColor(189, 183, 107),  # darkkhaki
    QColor(0, 128, 128),    # teal
    QColor(Qt.darkRed),
    QColor(0, 191, 255),    # deepSkyBlue
    QColor(216, 191, 216),  # thistle
    QColor(75, 0, 130),		# indigo
    QColor(255, 218, 185),  # peachPuff
    QColor(173, 255, 47),   # greenYellow
    QColor(0, 250, 154),    # mediumSpringGreen
    QColor(Qt.darkBlue),
    QColor(65, 105, 225),   # royalBlue
    QColor(176, 224, 230),  # powderBlue
    QColor(193, 255, 193),
    QColor(139, 102, 139),
    QColor(135, 206, 255),
    QColor(138, 43, 226),
    QColor(139, 10, 80)
]


class SignalPlot(pg.GraphicsLayoutWidget):
    def __init__(self, *args):
        super(SignalPlot, self).__init__(*args)

        self._plt = self.addPlot()
        self._plt.showGrid(x=True, y=True)
        self._plt.setMouseEnabled(x=True, y=False)
        self._plt.ctrl.fftCheck.setVisible(False)
        self._legend = self._plt.addLegend()
        self._chn_plt = dict()
        self._chn_data = dict()

        self._v_line = pg.InfiniteLine(angle=90, movable=False)
        self._h_line = pg.InfiniteLine(angle=0, movable=False)
        self._v_line.setVisible(False)
        self._h_line.setVisible(False)
        self._plt.addItem(self._v_line, ignoreBounds=True)
        self._plt.addItem(self._h_line, ignoreBounds=True)

        self._plt.setXRange(-1, 1)
        self._plt.setYRange(-1, 1)
        self._pos_label = pg.TextItem(text=" ")
        self._pos_label.setPos(0, 0)
        self._pos_label.setVisible(False)
        self._plt.addItem(self._pos_label)
        self._plt.getViewBox().addedItems.remove(self._pos_label)
        self.scene().sigMouseMoved.connect(self.mouse_moved)

    def chn_data(self):
        return self._chn_data.copy()

    def set_chn_data(self, chn: int, y=[],  x=None):
        if chn >= 32 or chn < 0:
            raise RuntimeError("chn {} invalid".format(chn))

        if len(x) != len(y):
            raise RuntimeError("data x len must equals y len")

        if len(self._chn_data) == 0:
            self._plt.enableAutoRange()

        self._chn_data[chn] = (x, y)

        if chn not in self._chn_plt:
            plt = self._plt.plot(pen=chn_color[chn], name="chn {}".format(chn))
            plt.setZValue(-(chn+1))
            self._chn_plt[chn] = plt
            plt.getViewBox().menu.viewAll.triggered.connect(lambda: self.update_scale_data(chn, True))

        self.update_scale_data(chn)

    def update_scale_data(self, chn, force_view_all=False):
        plt = self._chn_plt[chn]
        x, y = self._chn_data[chn]
        data_len = len(y)
        if data_len < 2:
            return
        if x is None:
            x = np.linspace(0, 1, data_len)

        if (force_view_all or plt.getViewBox().getState()["autoRange"][pg.ViewBox.XAxis] == 1) and len(x) >= 2:
            x_start, x_len = x[0], x[-1]-x[0]
        else:
            x_start, _, x_len, _ = plt.viewRect().getRect()

        data_x_range = x[-1] - x[0]
        start_idx = (x_start - 0.5 * x_len - x[0]) / (data_x_range / data_len)
        start_idx = 0 if start_idx < 0 else int(start_idx)
        stop_dix = (x_start + 1.5 * x_len - x[0]) / (data_x_range / data_len)
        stop_dix = data_len if stop_dix >= data_len else int(stop_dix)

        scale = int((stop_dix-start_idx) / self.viewRect().width() / 2)
        if scale == 0:
            scale = 1

        scaled_x, scaled_y = self.scale_data(x[start_idx: stop_dix], y[start_idx: stop_dix], scale)

        data = np.array([scaled_x, scaled_y]).transpose((1, 0))
        plt.setData(data)
        if force_view_all:
            plt.getViewBox().autoRange()

    @staticmethod
    def scale_data(x_data, y_data, scale):
        if len(x_data) != len(y_data):
            return [], []
        data_len = len(x_data)

        if scale <= 2:
            scaled_x = x_data
            scaled_y = y_data
        else:
            scaled_x = []
            scaled_y = []

            for i in np.arange(0, data_len-scale, scale):
                data = y_data[i:i + scale]
                min_idx = i + np.argmin(data)
                max_idx = i + np.argmax(data)
                if min_idx == max_idx:
                    scaled_x.append(x_data[min_idx])
                    scaled_y.append(y_data[max_idx])
                elif min_idx < max_idx:
                    scaled_x.extend([x_data[min_idx], x_data[max_idx]])
                    scaled_y.extend([y_data[min_idx], y_data[max_idx]])
                else:
                    scaled_x.extend([x_data[max_idx], x_data[min_idx]])
                    scaled_y.extend([y_data[max_idx], y_data[min_idx]])
        return np.array(scaled_x), np.array(scaled_y)

    def set_chn_visible(self, chn: int, visible: bool):
        if chn in self._chn_plt:
            self._chn_plt[chn].setVisible(visible)

    def mouse_moved(self, e):
        pos = self._plt.vb.mapSceneToView(e)
        if self._plt.getViewBox().sceneBoundingRect().contains(e):
            self._pos_label.setText("x: {:,.2f}, y: {:,.2f}".format(pos.x(), pos.y()))
            self._v_line.setPos(pos.x())
            self._h_line.setPos(pos.y())

        text_size = self._pos_label.boundingRect()
        view_bound = self._plt.getViewBox().sceneBoundingRect()
        x_pad = np.ceil(view_bound.width() * self._plt.getViewBox().suggestPadding(pg.ViewBox.XAxis))
        y_pad = np.ceil(view_bound.height() * self._plt.getViewBox().suggestPadding(pg.ViewBox.YAxis))
        view_bound = QRectF(view_bound.left()+x_pad, view_bound.top()+y_pad, view_bound.width()-2*x_pad, view_bound.height()-2*y_pad)

        if e.x() <= view_bound.left():
            x = view_bound.left()
        elif e.x() + text_size.width() <= view_bound.right():
            x = e.x()
        else:
            x = view_bound.right() - text_size.width()

        if e.y() <= view_bound.top():
            y = view_bound.top()
        elif e.y() + text_size.height() <= view_bound.bottom():
            y = e.y()
        else:
            y = view_bound.bottom() - text_size.height()
        pos = self._plt.mapToView(QPointF(x, y))
        self._pos_label.setPos(pos.x(), pos.y())

    def keyPressEvent(self, ev):
        if ev.key() == Qt.Key_Control:
            self._plt.setMouseEnabled(x=False, y=True)
        elif ev.key() == Qt.Key_Alt:
            self._plt.setMouseEnabled(x=True, y=True)
        super(SignalPlot, self).keyPressEvent(ev)

    def keyReleaseEvent(self, ev):
        if ev.key() == Qt.Key_Control or ev.key() == Qt.Key_Alt:
            self._plt.setMouseEnabled(x=True, y=False)
        super(SignalPlot, self).keyReleaseEvent(ev)

    def wheelEvent(self, ev):
        super(SignalPlot, self).wheelEvent(ev)
        for chn in self._chn_data:
            if self._chn_plt[chn].isVisible():
                self.update_scale_data(chn)

    def mouseMoveEvent(self, ev):
        if ev.buttons() & Qt.LeftButton:
            for chn in self._chn_data:
                if self._chn_plt[chn].isVisible():
                    self.update_scale_data(chn)
        super(SignalPlot, self).mouseMoveEvent(ev)

    def paintEvent(self, ev):
        if self._plt.mouseHovering:
            self._pos_label.setVisible(True)
            self._v_line.setVisible(True)
            self._h_line.setVisible(True)
        else:
            self._pos_label.setVisible(False)
            self._v_line.setVisible(False)
            self._h_line.setVisible(False)

        if self._plt.getViewBox().getState()["autoRange"][pg.ViewBox.XAxis] == 1:
            for chn, (x, y) in self._chn_data.items():
                if len(self._chn_plt[chn].getData()[0]) != len(x):
                    self.update_scale_data(chn, force_view_all=True)
        super(SignalPlot, self).paintEvent(ev)

 

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值