PYTHON- Removing colorbar from imbedded figure in tkinter when using xarray's methods to plot

Question

I am currently implementing a simple GUI used to display Netcdf content as well as a few other uses, however I am getting stuck on one question that appeared simpler than it was: How to remove the colorbar created by a xarray dataset?

What I want to do goes as follow:

• Create a Figure imbedded in a tkinter Frame and display some data in it
• Have the possibility to freely plot one variable or the other

Everything goes well until I start plotting 2D variables (see example bellow, variable d2), xarray will automatically create a colorbar to match the plot (which is the desired effect), however I don't know how to remove this colorbar when switching back to displaying another variable. This results in the colorbar still being in the figure and in a new colorbar being created next to the first one when we switch back to displaying the d1 in the example.

As such is there a way to either access the axis containing the colorbar created by xarray and remove it or to access the colorbar itself and remove it?

import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
import numpy as np
import tkinter as tk
import xarray as xr

class MainFrame(tk.Frame):
    def __init__(self, master, ds):
        # dummy data
        self.ds = ds
        self.plt_displayed = 'd1'
        
        self.parent = master
        tk.Frame.__init__(self, self.parent, bg='turquoise', borderwidth=1, relief="sunken")
        self.create_layout()
        
        # Protocol
        self.parent.protocol("WM_DELETE_WINDOW", self.window_closure_handler)
        
    def window_closure_handler(self):
        self.parent.quit()
        self.parent.destroy()

    def create_layout(self):
        # Setting up the button
        click_me = tk.Button(self, text='Click Me!', command=self.bind_button_click_me)
        
        # Seting up the ploting area
        self.make_background_figure()
        self.canvas = FigureCanvasTkAgg(self.fig, master=self)
        self.canvas.draw()
        
        toolbar = NavigationToolbar2Tk(self.canvas, self, pack_toolbar=True)
        toolbar.update()
        
        # Setting up the geometry
        toolbar.pack(side=tk.BOTTOM, fill=tk.X)
        plot_widget = self.canvas.get_tk_widget()
        plot_widget.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
        click_me.pack()
    
    def refresh_canvas(self):
        self.axs.cla()
        self.make_background_figure()
        self.canvas.draw()    
    
    def make_background_figure(self):
        if not hasattr(self, 'fig'):
            self.fig, self.axs = plt.subplots(figsize=(5,4), dpi=100)
        self.ds['d1'].plot(ax=self.axs)
        self.axs.set(title='This 1D data!')
        self.axs.grid()
        
    def plt_d2(self):
        self.ds['d2'].plot(ax=self.axs)
        self.axs.set(title='This 2D data!')
        self.canvas.draw()
        
    def bind_button_click_me(self):
        if self.plt_displayed == 'd1':
            self.plt_d2()
            self.plt_displayed = 'd2'
        else:
            self.refresh_canvas()
            self.plt_displayed = 'd1'
             
    
def main():
    # Creating some dummy data
    x = np.arange(0, 100, 0.1)
    y = np.flip(x)
    d1 = np.sin(2*np.pi*x/15)
    d2 = d1[:, None] * np.cos(2*np.pi*y/15)[None, :]
    
    ds = xr.Dataset(
        data_vars=dict(d1=(["x"], d1), d2=(["y", "x"], d2)),
        coords=dict(x=("x", x), y=("y", y))
        )
        
    root = tk.Tk()
    
    window = MainFrame(root, ds)
    window.pack(side='top', fill='both', expand=True)
    
    root.mainloop()
            

if __name__ == '__main__':
    main()

Additional notes:

• I know, it is possible to directly bypass xarray's plot methods and use matplotlib instead but xarray's ones are already implemented and let me provide a CF-compliant dataset to automatically manage the figure labels
• While it is possible to destroy the widget containing the figure and create a new one, I would like to know if it is possible to simply "refresh" the figure with a new plot

I am using:

• xarray -> 2024.3.0
• matplotlib -> 3.8.4
• tkinter -> 8.6.14

Answer 1

Thanks @Stef for your answer, it gave me a direction to look into. Using your answer I started looking into how to remove the axe containing the colorbar by calling it using fig.axes[-1] but while it can be done, it also leaves an empty white space where the colorbar container was.

Instead of blindly removing axes again and again, I came to the conclusion that it would be better to clear the whole figure and only re-define the axe I need. Bellow is the part of the original I changed to suit my purpose:

def make_background_figure(self):
    if not hasattr(self, 'fig'):
        self.fig, self.axs = plt.subplots(figsize=(5,4), dpi=100)
    
    elif len(self.fig.axes) > 1:
        self.fig.clf()
        self.axs = self.fig.add_subplot(1,1,1)
        
    self.ds['d1'].plot(ax=self.axs)
    self.axs.set(title='This 1D data!')
    self.axs.grid()

The only difference being the elif condition allowing me to check whether a colobar is present, in which case the figure is cleared and a new subplots is created to replace the one which just got removed by fig.clf()