I'm writing a simple function to plot 2D data. I want the colorbar label on the x-axis of the colorbar, vertically aligned with the xlabel of the plot. Here's a MWE:
from matplotlib import pyplot as plt
def plot2d(x, y, z, xlabel, ylabel, zlabel, ax=None, **pcm_kwargs):
if ax is None:
fig, ax = plt.subplots(constrained_layout=True)
else:
fig = ax.get_figure()
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
pcm_kwargs.setdefault("shading", "auto")
mesh = ax.pcolormesh(x, y, z, **pcm_kwargs)
cb = fig.colorbar(mesh, ax=ax)
cb.ax.set_xlabel(zlabel)
fig.align_xlabels()
return fig, ax
if __name__ == '__main__':
import numpy as np
x = np.linspace(0, 10, 100)
y = np.linspace(0, 10, 100)
z = np.random.random((100, 100))
plot2d(x, y, z, xlabel="xlabel", ylabel="ylabel", zlabel="zlabel")
plt.show()
First problem is that constrained_layout=True
causes fig.colorbar
to create an Axes
instead of an AxesSubplot
, so fig.align_xlabels()
throws an error. But even if I remove constrained_layout=True
, the xlabel and zlabel aren't vertically aligned:
.
I also tried setting the label's y position directly:
fig.canvas.draw()
cb.ax.xaxis.label.set_y(ax.xaxis.label.get_position()[1])
but still no effect.
How can I vertically align the plot and colorbar xlabels (within a constrained_layout
if possible)?