[PyQt5] 선택한 figure 창에 그래프 그려주는 프로그램

다음 글에서 이어지는 내용입니다.

[PyQt5] 새 창을 띄워서 matplotlib.animation으로 실시간 그래프 그리기, https://stella47.tistory.com/445

 

앞선 글에서 plot은 지정되어있는 데이터만을 출력했다.

본 글은 사용자가 특정 figure에 선택적으로 데이터를 출력하고 지우는 기능을 추가했다.

 

결과물

창을 조작하는 모습은 다음과 같다.

animate 함수에 데이터 업데이트가 같이 껴있어서 Pause를 누르면 데이터가 끊긴다.

 

---

코드 구현

1. pypubsub에 토픽 figure_cmd_rx을 통해서 어떤 데이터를 그리고 지울 지 전달한다,

2. FigureMaker는 사용자 명령에 따라서 데이터를 그려준다.

3. Tester는 콤보 박스를 통해서 어떤 Figure에 어떤 데이터를 선택할지를 고를 수 있다.

더보기를 누르면 코드를 볼 수 있다.

 

import sys
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtWidgets import QPushButton, QLabel, QComboBox
from PyQt5.QtWidgets import QMainWindow, QApplication, QWidget, QVBoxLayout
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.backend_bases import FigureCanvasBase
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
 
from pubsub import pub
import multiprocessing as mp
import time
import math
 
from stop_watch import StopWatch
 
QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
 
class FigureBase(QWidget):
    def __init__(self):
        super().__init__()
        
        # Default Figure Setting
        self._time_scale = 10
 
        pub.subscribe(self._cb_figure_cmd_rx, "figure_cmd_rx")
        pub.subscribe(self._cb_figure_data_rx, "figure_data_rx")
        # To check running time
        self._stop_watch_draw = StopWatch()
        self._stop_watch_module = StopWatch()
        
        self._init_draw()
        self._init_UI()
        self.setGeometry(310,300,400,600)
        self.show()
 
    def __del__(self):
        plt.close()
 
    # override figure close event from QWidget
    def closeEvent(self, event):
        plt.close()
        self.close()
 
    def _init_UI(self):
        vbox = QVBoxLayout()
        vbox.addWidget(self.toolbar)
        vbox.addWidget(self.canvas)
        self.setLayout(vbox)
        self.setGeometry(400,300,400,600)
 
    def _init_draw(self):
        # sub-class will store data at this lists
        self._number_of_figures = 3
        self._time = list()
        self._data = list()
        self._data.append(list([]))
        self._data.append(list([]))
        self._data.append(list([]))
        self.fig = plt.Figure()
        self.canvas = FigureCanvas(self.fig)
        self.toolbar = NavigationToolbar(self.canvas, self)
 
        # Figure Settings
        cfg = list()
        cfg.append(list([0]))
        cfg.append(list([]))
        cfg.append(list([]))
        # make figure list
        self._plots = list()
        for idx in range(0, self._number_of_figures):
            self._plots.append({"FIGURE":self.fig.add_subplot(self._number_of_figures,1,idx+1), 'INDEX':cfg[idx]})
            self._plots[idx]['FIGURE'].clear()
 
        # 애니메이션 생성하고 그리기
        self.ani = animation.FuncAnimation(fig=self.fig, func=self.animate, interval=10, blit=False, save_count=100)
        self.canvas.draw()
 
    # 그리기 전에 아무거나 계산
    def _before_drawing(self):
        self._stop_watch_module.click()
        pub.sendMessage("draw_dt",msgData=self._stop_watch_module.get_dt())
 
        x = self._stop_watch_draw.get_time_since()
 
        self._time.append(x)
        self._data[0].append(math.sin(x))
        self._data[1].append(math.cos(x))
        self._data[2].append(math.cos(x)*2.0)
 
    # Sub-class must override this function
    def animate(self, event):
        self._before_drawing()
 
        min_idx = self._find_x_lim()
        # min_idx = 0
        for item in self._plots:
            item['FIGURE'].clear()
            item['FIGURE'].grid()
            if(len(self._time) >= 2):
                minmax = [self._time[-1]-self._time_scale, self._time[-1]]
                item['FIGURE'].set_xlim(minmax[0], minmax[1])
            # Put each data @ sub-class
            for idx in item['INDEX']:
                item['FIGURE'].plot(self._time, self._data[idx])
                # item['FIGURE'].plot(self._time[-min_idx:], self._data[idx][-min_idx:])
 
        self._stop_watch_module.click()
        pub.sendMessage("draw_dt",msgData=self._stop_watch_module.get_dt())
 
    def _cb_figure_cmd_rx(self, msgData):
        # print(msgData)
        figure  = msgData['FIGURE']
        idx     = msgData['INDEX']
        cmd     = msgData['COMMAND']
        if(cmd == 'ADD'):
            if(idx not in self._plots[figure]['INDEX']):
                self._plots[figure]['INDEX'].append(idx)
        else: # REMOVE
            if(idx in self._plots[figure]['INDEX']):
                self._plots[figure]['INDEX'].remove(idx)
        print(self._plots)
        # pass
 
    def _cb_figure_data_rx(self):
        pass
 
    # 보여줄 크기에 맞는 인덱스를 찾는다.
    def _find_x_lim(self):
        min_idx = 0
        for idx in range(len(self._time)):
            if(self._time[-1] - self._time[-idx-1] > self._time_scale):
                min_idx = idx+1
                break
        return min_idx
    
    def pause(self):
        # stop 은 멈추는데
        # start 가 아예 새로 시작함
        # self.ani.event_source.stop()
        self.ani.pause()
 
    def resume(self):
        # start나 resume이랑 다를 바 없이 다 지우고 시작한다.
        # self.ani.event_source.start()
        self.ani.resume()
 
    def close(self):
        self.ani.event_source.stop() # 애니메이션 정지
        # self._timer.stop() # 데이터 업데이트 정지
        # plt.close() # 이거 해봤자 창이 안꺼짐.
        super().close() # 이걸 해야 위젯이 꺼진다.
        
class Tester(QMainWindow):
 
    def __init__(self):
        super().__init__()
        self._widget = None  # No external window yet.
 
 
        self.cb = list()
        cb = QComboBox(self)
        cb.setGeometry(10, 90, 90, 30)
        cb.addItem("Fig.1")
        cb.addItem("Fig.2")
        cb.addItem("Fig.3")
        self.cb.append(cb)
        cb = QComboBox(self)
        cb.setGeometry(110, 90, 90, 30)
        cb.addItem("data.1")
        cb.addItem("data.2")
        cb.addItem("data.3")
        self.cb.append(cb)
 
        self.button = list()
        button = self._get_button('Create', parent=self, geo=[10, 10, 90, 30],  callback=self._create_plot)
        self.button.append(button)
        button = self._get_button('Close',  parent=self, geo=[10, 50, 90, 30],  callback=self._close_plot)
        self.button.append(button)
        button = self._get_button('Pause',  parent=self, geo=[110, 10, 90, 30], callback=self._pause_plot)
        self.button.append(button)
        button = self._get_button('Resume', parent=self, geo=[110, 50, 90, 30], callback=self._resume_plot)
        self.button.append(button)
        button = self._get_button('Add',    parent=self, geo=[10, 130, 90, 30], callback=self._add_plot)
        self.button.append(button)
        button = self._get_button('Remove', parent=self, geo=[110, 130, 90, 30], callback=self._remove_plot)
        self.button.append(button)
 
        self.label = QLabel("dt :   ms", self)
        self.label.setGeometry(10, 170, 120, 30)
        font = self.label.font()
        font.setPointSize(10)
        self.label.setFont(font)
 
        pub.subscribe(self._cb_dt, 'draw_dt')
        self.dt = 0
 
        self.setGeometry(100,300,210,200)
        self.show()
 
    def _create_plot(self):
        if self._widget is not None:
            if (self._widget.isVisible() == False): # 우측 상단 X으로 끄면 visible이 False가 된다.
                self._close_plot()
            else:
                return
        self._widget = FigureBase() # 아무튼 초기화 했으니 켠다.
        # proc = mp.Process(target=self._widget.animate, )
        # proc.start()
        # proc.join()
 
    def _pause_plot(self):
        if self._widget is not None:
            self._widget.pause()
 
    def _resume_plot(self):
        if self._widget is not None:
            self._widget.resume()
 
    def _close_plot(self):
        if self._widget is not None:
            self._widget.close()
            self._widget.__del__()
            self._widget = None
 
    def _add_plot(self):
        pub.sendMessage("figure_cmd_rx", msgData={"FIGURE":self.cb[0].currentIndex(), "INDEX":self.cb[1].currentIndex(), "COMMAND":"ADD"})
        # print("Add {} to {}".format(self.cb[1].currentText(), self.cb[0].currentText()))
 
    def _remove_plot(self):
        pub.sendMessage("figure_cmd_rx", msgData={"FIGURE":self.cb[0].currentIndex(), "INDEX":self.cb[1].currentIndex(), "COMMAND":"REMOVE"})
        # print("Remove {} to {}".format(self.cb[1].currentText(), self.cb[0].currentText()))
 
    def _get_button(self, name, parent, geo, callback=None):
        button = QPushButton(name, parent)
        button.setGeometry(geo[0],geo[1],geo[2],geo[3])
        if(callback != None): button.clicked.connect(callback)
        return button
 
    def _cb_dt(self, msgData):
        self.dt = msgData
        self.label.setText("dt : {:10.3f} ms".format(self.dt))
        
if __name__ == "__main__":
 
    print('support_blit : ', FigureCanvasBase.supports_blit)
    app = QApplication(sys.argv)
    tester = Tester()
    app.exec()

 

EOF