How to programtically set a suitable padding for matplotlib colorbar?

I’m creating some subplots each with its own colorbar at the bottom. The colorbar is added using:

cax, kw = mcbar.make_axes_gridspec(ax, orientation='horizontal',
                                   pad=pad,
                                   fraction=0.07, shrink=0.85, aspect=35)
figure.colorbar(cs, cax=cax, orientation='horizontal')

The pad argument is adjusted, so that if there is no xticklabels, the value is smaller, to avoid wasting space.

The complete script:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colorbar as mcbar

x = np.linspace(-1, 1, 100)
y = np.linspace(-1, 1, 50)
X, Y = np.meshgrid(x, y)
Z = X**2+np.sin(Y)

figure = plt.figure(figsize=(12, 10))
nrow = 3
ncol = 2

for ii in range(nrow*ncol):
    ax = figure.add_subplot(nrow, ncol, ii+1)
    row, col = np.unravel_index(ii, (nrow, ncol))

    cs = ax.contourf(X, Y, Z)
    if row == nrow-1:
        # larger padding to make room for xticklabels
        pad = 0.15
    else:
        # smaller padding otherwise
        pad = 0.05
        ax.tick_params(labelbottom=False)

    if row == 1 and col == 1:
        # add xlabel would need more padding
        ax.set_xlabel('X')

    cax, kw = mcbar.make_axes_gridspec(ax, orientation='horizontal',
                                       pad=pad,
                                       fraction=0.07, shrink=0.85, aspect=35)
    figure.colorbar(cs, cax=cax, orientation='horizontal')
    ax.set_title(str(ii+1))

figure.tight_layout()
figure.show()

The output figure:

enter image description here

But the current solution is using hard-coded padding values (0.15 if with xticklabels, 0.05 otherwise), and it doesn’t adjust well to the existence of xlabels (see subplot 4), or changing figure sizes.

Is there a way to programmatically work out a suitable padding value to place the colorbar? Maybe by adjusting the bounding box of the parent axis object so that its bbox is smaller if there is no xlabels or xticklabels, or by finding out the coordinates of the parent axis and somehow computing a padding?

Answer

You can get the space needed for tick labels and the axis label by comparing the bounding boxes of the whole axes and the yaxis. To get these bounding boxes we need a renderer. To make it available we first need to draw the canvas. The bounding boxes are returned in display coordinates, so we transform them to axes coordinates using the inverted axes transformation. The difference of their y coordinates gives the required extra padding:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colorbar as mcbar
from matplotlib.transforms import Bbox

x = np.linspace(-1, 1, 100)
y = np.linspace(-1, 1, 50)
X, Y = np.meshgrid(x, y)
Z = X**2+np.sin(Y)

figure = plt.figure(figsize=(12, 10))
figure.canvas.draw()  # to get renderer
nrow = 3
ncol = 2

for ii in range(nrow*ncol):
    ax = figure.add_subplot(nrow, ncol, ii+1)
    row, col = np.unravel_index(ii, (nrow, ncol))

    cs = ax.contourf(X, Y, Z)
    if row != nrow-1:
        ax.tick_params(labelbottom=False)

    if row == 1 and col == 1:
        # add xlabel would need more padding
        ax.set_xlabel('X')

    # get height of ticklabels and label
    b = ax.transAxes.inverted().transform(
        [ax.yaxis.get_tightbbox(figure.canvas.renderer).p0,
        ax.get_tightbbox(figure.canvas.renderer).p0]
    )
    pad = 0.05 + (b[0]-b[1])[1]

    cax, kw = mcbar.make_axes_gridspec(ax, orientation='horizontal',
                                       pad=pad,
                                       fraction=0.07, shrink=0.85, aspect=35)
    figure.colorbar(cs, cax=cax, orientation='horizontal')
    ax.set_title(str(ii+1))

enter image description here

This solution has the flaw that axes 3 and 4 have different heights. You can fix this by adjusting ymin of all axes in a row to the row maximum:

figure.tight_layout()
for i in range(0, 2*ncol*nrow, 2*ncol):
    ymin = 0
    for j in range(0, 2*ncol, 2):
        ymin = max(ymin, figure.axes[i+j].get_position().ymin)
    for j in range(0, 2*ncol, 2):
        b = figure.axes[i+j].get_position()
        figure.axes[i+j].set_position(Bbox([[b.xmin,ymin],[b.xmax,b.ymax]]))

Please note that this adjustment must be done before applying tight_layout! enter image description here