Search code examples
pythonpython-3.xmatplotlibpyqt5colorbar

How to remove Matplotlib colorbar upon subsequent replots?


I'm trying to do something pretty straight forward in Matplotlib, but hitting a ton of hurdles; after googling for hours, I'm thinking this takes some very specific magic. I just want my GUI to start on a blank screen and then import a file that needs a colorbar - if the file is good, a colorbar is added. If the file is added again (or a different one), the colorbar is removed and a new one is plotted. If the file is bad, we reset to initial conditions (no colorbars or plot). I really hope the below code makes sense for what I'm getting at here, thank you so much for your help and time in advance:

import sys
import os
import random
import matplotlib
matplotlib.use('Qt5Agg')
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtWidgets import QGridLayout, QFileDialog, QPushButton

import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure


class MyMplCanvas(FigureCanvas):

    def __init__(self, parent=None, width=5, height=4, dpi=100):
        MyMplCanvas.fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = MyMplCanvas.fig.add_subplot(111)

        self.compute_initial_figure()

        FigureCanvas.__init__(self, MyMplCanvas.fig)
        self.setParent(parent)

        FigureCanvas.setSizePolicy(self,
                                   QtWidgets.QSizePolicy.Expanding,
                                   QtWidgets.QSizePolicy.Expanding)
        FigureCanvas.updateGeometry(self)

    def compute_initial_figure(self):
        pass


class MyDynamicMplCanvas(MyMplCanvas):

    def __init__(self, *args, **kwargs):
        MyMplCanvas.__init__(self, *args, **kwargs)

    def compute_initial_figure(self):
        C=np.random.rand(500).reshape((20,25))
        S=np.random.rand(500).reshape((20,25))
        dc = self.function(S,C)
        im = self.axes.imshow(dc, alpha = 0)

        MyMplCanvas.fig.patch.set_facecolor('k')
        self.axes.patch.set_facecolor('k')

        # don't want to show the cb initially (imagine GUI starts on a blank screen)
        #cb = MyMplCanvas.fig.colorbar(im, ax=self.axes,
                                    #orientation='vertical')#,use_gridspec=True)??

    def update_figure(self):
        # need something like the below to delete subsequent colorbars if they exist:
        #if MyMplCanvas.fig.axes[1] is not None:
#            MyMplCanvas.fig.delaxes(MyMplCanvas.fig.axes[1]) #remove colorbar
        self.axes.cla()

        if P1.df: #datafile is good
            MyMplCanvas.fig.patch.set_facecolor('w')
            self.axes.patch.set_facecolor('w')
            C=np.random.rand(500).reshape((20,25))
            S=np.random.rand(500).reshape((20,25))
            dc = self.function(S,C)

            im = self.axes.imshow(dc)
            cb = MyMplCanvas.fig.colorbar(im, ax=self.axes,
                                    orientation='vertical')#,use_gridspec=True)??
        else: #datafile was bad, or they pressed cancel
            C=np.random.rand(500).reshape((20,25))
            S=np.random.rand(500).reshape((20,25))
            dc = self.function(S,C)
            im = self.axes.imshow(dc, alpha = 0)

            MyMplCanvas.fig.patch.set_facecolor('k')
            self.axes.patch.set_facecolor('k')

        self.show()
        self.draw()

    def function(self,s,c):
        return s*2+c

class P1(QtWidgets.QWidget):

    df = False

    def __init__(self, parent=None):
        super(P1, self).__init__(parent)
        layout = QGridLayout(self)

        self.button_browse1 = QPushButton('Browse Good', self)
        self.button_browse1.clicked.connect(self.browseFile1)
        layout.addWidget(self.button_browse1, 1, 1, 1, 1)
        self.button_browse1.show()

        self.button_browse2 = QPushButton('Browse Bad', self)
        self.button_browse2.clicked.connect(self.browseFile2)
        layout.addWidget(self.button_browse2, 2, 1, 1, 1)
        self.button_browse2.show()

        self.dc = MyDynamicMplCanvas(self, width=5, height=4, dpi=100)
        layout.addWidget(self.dc, 3, 1, 1, 1)

    def browseFile1(self):
        P1.df = True
        self.dc.update_figure()

    def browseFile2(self):
        P1.df = False
        self.dc.update_figure()

class MainWindow(QtWidgets.QMainWindow):    
    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)
        self.setAttribute(QtCore.Qt.WA_DeleteOnClose)

        self.stack = QtWidgets.QStackedWidget(self)
        P1f = P1(self)
        self.stack.addWidget(P1f)
        self.setCentralWidget(self.stack)

if __name__ == '__main__':
    qApp = QtWidgets.QApplication(sys.argv)
    aw = MainWindow()
    aw.show()
    sys.exit(qApp.exec_())

Solution

  • Okay, so for anyone else stumbling on this, the trick is to first define your colorbar the very first time you plot (initial figure) via the top Class its in, i.e., MyMplCanvas.cbar = MyMplCanvas.fig.colorbar(im). Its very important to define it up front even if you didn't want to even use it yet.

    Then at the top of the update plot function, simply call MyMplCanvas.cbar.remove().