pythonmatplotlibwidgetcolorbar

Problems with updating colorbar with matplotlib.Slider listener in Jupyter

The problem that I am having is during updating the colorbar of an plt.imshow plot. Here's the code I will try to break it down and to explain some stuff in it afterwards.

# Imports ##############################################################
import numpy as np
%matplotlib notebook
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.widgets import Slider
from scipy.ndimage import label, find_objects
# Functions ############################################################
def intensity_distribution(r, z, PMax, w0, zR):
    wZ = w0 * np.sqrt(1 + (z / zR)**2)  # beam radius at z (Gaussian beam spreading)
    I0 = 2 * PMax / (np.pi * wZ**2) # peak intensity at radius wZ
    return I0 * np.exp(-2 * r**2 / wZ**2), wZ  # Gaussian intensity distribution, beam radius
def get_circle_ROI(I, r, threshold):
    ROI = (I > threshold).astype("uint8")  # binary mask for regions above the threshold
    labels, features = label(ROI) # label connected regions
    slices = find_objects(labels) # get bounding box slices for labeled regions
    xSlice, ySlice = slices[0] # extract x and y slices of the largest feature
    ROIHeight = (xSlice.stop - xSlice.start) * (r[1] - r[0]) * 1e6 # convert height to micrometers
    ROIWidth = (ySlice.stop - ySlice.start) * (r[1] - r[0]) * 1e6 # convert width to micrometers
    cx = (ySlice.start + ySlice.stop) // 2 # x-coordinate of the center
    cy = (xSlice.start + xSlice.stop) // 2 # y-coordinate of the center
    centre = (r[cy] * 1e6, r[cx] * 1e6) # convert center coordinates to micrometers
    radius = min(ROIWidth, ROIHeight) / 2 # radius is the smaller dimension's half-width
    return centre, radius
def update_plot(PMax, zOffset):
    """Update the heatmap based on new parameters."""
    global colorbar
    # Calculate intensity distribution at given z offset
    I, wZ = intensity_distribution(np.sqrt(R**2 + Z**2), zOffset, PMax, BEAM_RADIUS_AT_FOCUS, zR)
    I /= 1e6  # convert intensity from W/m² to W/mm²
    I += 0.01  # small offset for better visualization contrast
    max_intensity = I.max()  # maximum intensity in the current distribution
    # Calculate the on-axis peak intensity at focus in W/mm²
    I0 = (2 * PMax) / (np.pi * BEAM_RADIUS_AT_FOCUS**2)  # peak intensity in W/m² at z = 0
    I0 /= 1e6  # convert peak intensity to W/mm²
    # Calculate the Full Width at Half Maximum (FWHM) in micrometers
    centre, fwhm = get_circle_ROI(I, r, max_intensity / 2)
    _, tenth = get_circle_ROI(I, r, max_intensity / 10)
    # Clear and update plot
    ax.clear()  # clear current axes
    # Display the updated intensity distribution as a heatmap
    im = ax.imshow(I, extent=[r[0]*1e6, r[-1]*1e6, r[0]*1e6, r[-1]*1e6], norm=colors.LogNorm(vmax=14000))
    ax.set_xlabel("x (μm)")  # label for x-axis
    ax.set_ylabel("y (μm)")  # label for y-axis
    # Add plot title with z offset, FWHM, and max intensity in W/mm²
    ax.set_title(f"FWHM = {fwhm:.1f} μm\n"
                 f"Radius at 10% of total power = {tenth:.2f} μm\n"
                 f"Max power = {I.max():.2f} W/mm²",
                 loc="left")
    # Draw a circle representing the FWHM boundary
    cirlcefwhm = plt.Circle(centre, fwhm, color='white', fill=False, linestyle='--', linewidth=2, label="FWHM")
    cirlce10 = plt.Circle(centre, tenth, color='white', fill=False, linestyle='--', linewidth=2, label="10% of I$_max$")
    ax.add_patch(cirlcefwhm)  # add the FWHM circle to the plot
    ax.add_patch(cirlce10)  # add the circle where power is 10% of max
    #### Problematic starts here ####
    if colorbar is not None:   # if colorbar already exists, remove it
        colorbar.remove()
    colorbar = plt.colorbar(im, ax=ax, label="Intensity (W/mm²)")  # create new colorbar in W/mm²
    fig.draw_without_rendering() # redraw based on the recommendation of matplotlib instead of colorbar.draw_all()
    fig.canvas.draw()  # redraw figure to reflect updates
    #### Problematic ends here ####
def sliders_on_changed(val):
    ''' Slider update function '''
    power = power_slider.val * MAX_LASER_POWER / 100  # calculate current power level in watts
    z_offset = z_offset_slider.val / 1000 # convert slider z offset from mm to meters
    update_plot(power, z_offset) # update the plot with new parameters
# Inputs ###############################################################
WAVELENGTH = 10.6e-6  # wavelength in meters
MAX_LASER_POWER = 80  # max laser power in watts
BEAM_WIDTH_AT_FOCUS = 120e-6  # beam width at focus in meters
BEAM_RADIUS_AT_FOCUS = BEAM_WIDTH_AT_FOCUS / 2  # beam radius at focus in meters
zR = np.pi * BEAM_RADIUS_AT_FOCUS**2 / WAVELENGTH  # Rayleigh range in meters
gridSize = 100 # resolution
r = np.linspace(-500e-6, 500e-6, gridSize)  # range for spatial coordinates in meters
R, Z = np.meshgrid(r, r) # create grid for spatial coordinates
colorbar = None # init colorbar

The problematic part would be this one:

#### Problematic starts here ####
if colorbar is not None:   # if colorbar already exists, remove it
    colorbar.remove()
colorbar = plt.colorbar(im, ax=ax, label="Intensity (W/mm²)")  # create new colorbar in W/mm²
fig.draw_without_rendering() # redraw based on the recommendation of matplotlib instead of colorbar.draw_all()
fig.canvas.draw()  # redraw figure to reflect updates
#### Problematic ends here ####

Am not so sure why it is not working. What happens is that the colorbar appears for the first update and then dissapears after a change in the Z-offset or the power percentage. Basically after a change in any of the values.

fig, ax = plt.subplots()
plt.subplots_adjust(left=0.25, bottom=0.35)  # leave space for sliders
update_plot(50, 0)
ax_power = plt.axes([0.25, 0.2, 0.65, 0.03], facecolor = "lightgray")
ax_z_offset = plt.axes([0.25, 0.15, 0.65, 0.03], facecolor = "lightgray")
power_slider = Slider(ax_power, 'Power (%)', 0.1, 100, valinit=50)
z_offset_slider = Slider(ax_z_offset, 'Z-Offset (mm)', 0, 5.0, valinit=0)
power_slider.on_changed(sliders_on_changed)
z_offset_slider.on_changed(sliders_on_changed)

Running this will get you to a UI that looks like this:

First update

And then when I move any of the sliders the colorbar disappears:

After moving slider

The updating of the colorbar seems to be a common problem, I saw a lot of stackoverflow questions about it, I even answered one once.

The solution however is not really working with the slider listeners.

P.S.: Changing the backend does not really help. However, using %matplotlib qt revealed a traceback:

Traceback (most recent call last):
  File "C:\ProgramData\anaconda3\Lib\site-packages\matplotlib\cbook\__init__.py", line 309, in process
    func(*args, **kwargs)
  File "C:\ProgramData\anaconda3\Lib\site-packages\matplotlib\widgets.py", line 603, in <lambda>
    return self._observers.connect('changed', lambda val: func(val))
                                                          ^^^^^^^^^
  File "C:\Users\User\AppData\Local\Temp\ipykernel_26252\2170461726.py", line 66, in sliders_on_changed
    update_plot(power, z_offset) # update the plot with new parameters
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\User\AppData\Local\Temp\ipykernel_26252\2170461726.py", line 57, in update_plot
    colorbar.remove()
  File "C:\ProgramData\anaconda3\Lib\site-packages\matplotlib\colorbar.py", line 1041, in remove
    self.ax.remove()
  File "C:\ProgramData\anaconda3\Lib\site-packages\matplotlib\artist.py", line 242, in remove
    self._remove_method(self)
  File "C:\ProgramData\anaconda3\Lib\site-packages\matplotlib\figure.py", line 944, in delaxes
    self._axstack.remove(ax)
  File "C:\ProgramData\anaconda3\Lib\site-packages\matplotlib\figure.py", line 92, in remove
    self._axes.pop(a)
KeyError: <Axes: label='<colorbar>', ylabel='Intensity (W/mm²)'>

Some stuff that I tried:

Can someone please help me in figuring out what's wrong? I understand ofc that the code is quite long, I can shorten it to a simple example if necessary.

I am using

Solution

  • I managed to solve this problem on my old version by using the following lines to update the colormap:

    if colorbar is not None: # if colorbar already exists, update it
    colorbar.mappable.set_clim(vmax=max_intensity)
else:    
    colorbar = plt.colorbar(im, ax=ax, label="Intensity (W/mm²)") # create new colorbar in W/mm²

    Updating Anaconda was not needed in this case. The code above works both for the old version and the old version (under the %matplotlib qt backend)

    Proof:

    Proof

    The advice about using ipympl lead to a RuntimeError under a clean new anaconda installation. So, really no need for that.

    Using interact() does not solve anything.