[PyQt5] MatplotLib의 차트를 Widget으로 사용하기

PyQt에서 차트를 위젯으로 사용하기 위해 코드입니다. 먼저 실행 결과는 다음과 같습니다.

하단에 컴보박스를 통해 2가지 항목을 선택할 수 있는데요. 항목을 선택할때마다 해당되는 항목의 그래프가 상단에 표시됩니다.

먼저 이를 위한 UI 구성을 위해 아래와 같은 코드를 작성합니다.

import sys
import numpy as np
from PyQt5.QtWidgets import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

class MyWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()

        self.setLayout(self.layout)
        self.setGeometry(200, 200, 800, 600)

    def initUI(self):
        self.fig = plt.Figure()
        self.canvas = FigureCanvas(self.fig)
        
        layout = QVBoxLayout()
        layout.addWidget(self.canvas)

        cb = QComboBox()
        cb.addItem('Graph1')
        cb.addItem('Graph2')
        cb.activated[str].connect(self.onComboBoxChanged)
        layout.addWidget(cb)

        self.layout = layout

        self.onComboBoxChanged(cb.currentText())

    def onComboBoxChanged(self, text):
        if text == 'Graph1':
            self.doGraph1()
        elif text == 'Graph2':
            self.doGraph2()

    def doGraph1(self):
        ....

    def doGraph2(self):
        ....
            
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MyWindow()
    window.show()
    app.exec_()

Qt5에서 Matplot을 사용하기 위해서는 matplotlib.backends.backend_qt5agg 패이지의 FigureCanvasQTAgg라는 이름의 위젯 클래스가 필요하다는 것을 알 수 있습니다. MyWindow 클래스의 doGraph1과 doGraph2 함수의 코드가 핵심인데, 그 내용은 아래와 같습니다.

   def doGraph1(self):
        x = np.arange(0, 10, 0.5)
        y1 = np.sin(x)
        y2 = np.cos(x)
        
        self.fig.clear()

        ax = self.fig.add_subplot(111)
        ax.plot(x, y1, label="sin(x)")
        ax.plot(x, y2, label="cos(x)", linestyle="--")
        
        ax.set_xlabel("x")
        ax.set_xlabel("y")
        
        ax.set_title("sin & cos")
        ax.legend()
        
        self.canvas.draw()

    def doGraph2(self):
        X = np.arange(-5, 5, 0.25)
        Y = np.arange(-5, 5, 0.25)
        X, Y = np.meshgrid(X, Y)
        Z = X**2 + Y**2
        
        self.fig.clear()
        
        ax = self.fig.gca(projection='3d')
        ax.plot_wireframe(X, Y, Z, color='black')

        self.canvas.draw() 

doGraph2 함수에 대한 실행 결과는 다음과 같습니다. 물론 이 함수의 실행은 컴보박스의 항목 중 Graph2를 선택했을때 실행됩니다.

참고로 이 글의 차트는 아래의 글을 참고로 하여 작성하였습니다.

파이썬의 matplotlib 노트