gettting a colorbar programmatically from an axis object

NNN 0
Viewed 26

Consider the following code in which data is being plotted within a loop. In order to not plot multiple colorbars, I remove them before plotting new ones. However, I am tracking these colorbars manually. Is there a way to get a reference to them using, say, something like ax[0].get_colorbar. That would make my code a lot simpler.

import matplotlib.pyplot as plt
import numpy as np

# How to get a colorbar from an axis?

nrows = 1
ncols = 2
nstep = 5

fig, ax = plt.subplots(nrows=nrows,ncols=ncols)

cb0 = None
cb1 = None

for istep in range(nstep):
    data = np.random.random(size=(5,5))
    imu0 = ax[0].pcolormesh(data)
    imu1 = ax[1].pcolormesh(data)
    
    # this code is for removing previously drawn colorbars
    # I would like to get a reference to the colorbar cb0 from ax0 
    # and then remove it. 
    # I do not want to keep track of the colorbars manually 
    
    if cb0 is not None:
        cb0.remove()
    if cb1 is not None:
        cb1.remove()
    cb0 = plt.colorbar(imu0,ax=ax[0])
    cb1 = plt.colorbar(imu1,ax=ax[1])
colorbar subplot matplotlib python
1 Answers

The following seems to do what I want:

import matplotlib.pyplot as plt
import numpy as np

# How to get a colorbar from an axis?

nrows = 1
ncols = 2
nstep = 10

fig, ax = plt.subplots(nrows=nrows,ncols=ncols)

cb0 = None
cb1 = None

for istep in range(nstep):
    data = np.random.random(size=(5,5))+istep
    imu0 = ax[0].pcolormesh(data)
    imu1 = ax[1].pcolormesh(data)

    if ax[0].collections[0].colorbar is None:
        cb0 = plt.colorbar(imu0,ax=ax[0])
    else:
        ax[0].collections[0].colorbar.update_normal(imu0)
        
    if ax[1].collections[0].colorbar is None:
        cb1 = plt.colorbar(imu1,ax=ax[1])
    else:
        ax[1].collections[0].colorbar.update_normal(imu1)
NNN0
Related