Matplotlib在subplot2grid布局下外置colorbar避免遮挡与轴不对齐问题

需求与初始效果 #

需要在多个子图中绘制两个高斯分布，首次实现效果如下：

初始实现代码 #

import numpy as np
import numpy.random as rnd
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

mu_x1 = 3
sig_x1 = 2

mu_x2 = 3
sig_x2 = 0.5

N = int(1e3)

gaussian1 = rnd.normal(mu_x1, sig_x1, N)
gaussian2 = rnd.normal(mu_x2, sig_x2, N)
fig = plt.figure()

ax1 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 0),
                       rowspan = 2,
                       colspan = 2,
                       fig = fig
                       )
hist2d = ax1.hist2d(gaussian1,
                    gaussian2,
                    bins = 100,
                    density = True,
                    cmap=plt.cm.jet
                    )
fig.colorbar(hist2d[3], ax=ax1)
ax1.set_xlabel("$x_1$")
ax1.set_ylabel("$x_2$")

ax2 = plt.subplot2grid(shape = (3, 3),
                       loc = (0, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )
ax2.hist(gaussian1,
         bins = 100
         )
ax2.xaxis.set_ticks_position('top')
ax2.set_xlabel("$x_1$")
ax2.xaxis.set_label_position('top')
ax2.set_ylabel("Cuentas")

ax3 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 2),
                       rowspan = 2,
                       colspan = 1,
                       fig = fig
                       )
ax3.hist(gaussian2,
         bins = 100,
         orientation=u'horizontal')
ax3.yaxis.set_ticks_position('right')
ax3.yaxis.set_label_position('right')
ax3.set_xlabel("Cuentas")
ax3.set_ylabel("$x_2$")

plt.show()

初始效果的问题：二维直方图子图的x轴和上方子图的x轴宽度不一致，无法完全对齐，根源是colorbar挤占了ax1的宽度。

调整后效果与新问题 #

将colorbar移到二维直方图下方后效果如下：

调整后代码 #

fig = plt.figure()

ax1 = plt.subplot2grid(shape = (4, 3),
                       loc = (1, 0),
                       rowspan = 2,
                       colspan = 2,
                       fig = fig
                       )
hist2d = ax1.hist2d(gaussian1,
                    gaussian2,
                    bins = 100,
                    density = True,
                    cmap=plt.cm.jet
                    )
#fig.colorbar(hist2d[3], ax=ax1)
ax4 = plt.subplot2grid(shape = (4, 3),
                       loc = (3, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )    
fig.colorbar(hist2d[3], cax=ax4, orientation="horizontal")
ax1.set_xlabel("$x_1$")
ax1.set_ylabel("$x_2$")
# 注：代码中n1、n2为标准差倍数，可自行定义，如n1=n2=1对应1σ区间
ax1.axhline(mu_x2 + sig_x2*n2,c="k")
ax1.axhline(mu_x2 - sig_x2*n2,c="k")
ax1.axvline(mu_x1 + sig_x1*n1,c="k")
ax1.axvline(mu_x1 - sig_x1*n1,c="k")

ax2 = plt.subplot2grid(shape = (4, 3),
                       loc = (0, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )
ax2.hist(gaussian1,
         bins = 100
         )
ax2.xaxis.set_ticks_position('top')
ax2.set_xlabel("$x_1$")
ax2.xaxis.set_label_position('top')
ax2.set_ylabel("Cuentas")
ax2.axvline(mu_x1 + sig_x1*n1,c="k")
ax2.axvline(mu_x1 - sig_x1*n1,c="k")
#ax2.set_xticks([])

ax3 = plt.subplot2grid(shape = (4, 3),
                       loc = (1, 2),
                       rowspan = 2,
                       colspan = 1,
                       fig = fig
                       )
ax3.hist(gaussian2,
         bins = 100,
         orientation=u'horizontal')
ax3.yaxis.set_ticks_position('right')
ax3.yaxis.set_label_position('right')
ax3.set_xlabel("Cuentas")
ax3.set_ylabel("$x_2$")
ax3.axhline(mu_x2 + sig_x2*n2,c="k")
ax3.axhline(mu_x2 - sig_x2*n2,c="k")

调整后的新问题：colorbar占满了分配的子图区域，尺寸过大，遮挡了上方子图的刻度。

推荐两种成熟方案，第一种无需修改原有子图布局，实现更简单：

方案1：用make_axes_locatable给二维直方图轴额外附加colorbar轴 #

你已经导入了make_axes_locatable工具，直接用它在ax1的底部分出一个小轴放colorbar，完全不破坏原有3x3的子图对齐关系，修改ax1部分的代码即可：

ax1 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 0),
                       rowspan = 2,
                       colspan = 2,
                       fig = fig
                       )
hist2d = ax1.hist2d(gaussian1,
                    gaussian2,
                    bins = 100,
                    density = True,
                    cmap=plt.cm.jet
                    )
# 新增以下代码
divider = make_axes_locatable(ax1)
# 底部添加高度为ax1的5%、间距0.1的轴放colorbar
cax = divider.append_axes("bottom", size="5%", pad=0.1)
fig.colorbar(hist2d[3], cax=cax, orientation="horizontal")
# 将ax1的x轴刻度和标签移到顶部，避免和colorbar重叠
ax1.xaxis.set_label_position('top')
ax1.tick_params(axis='x', labeltop=True, labelbottom=False)

修改后ax1和上方ax2的x轴会完全对齐，colorbar尺寸合适，不会遮挡任何内容。

方案2：调整4x3子图的行高比例 #

如果一定要保留4x3的子图布局，可通过gridspec设置各行的高度占比，缩小colorbar所在行的高度：

# 替换原有的fig = plt.figure()和subplot2grid创建逻辑
fig = plt.figure(constrained_layout=True)
# 定义4行3列网格，最后一行高度仅为0.3份，其他行按需求设置比例
gs = fig.add_gridspec(4, 3, height_ratios=[1, 2, 2, 0.3])
ax1 = fig.add_subplot(gs[1:3, 0:2])
ax2 = fig.add_subplot(gs[0, 0:2])
ax3 = fig.add_subplot(gs[1:3, 2])
ax4 = fig.add_subplot(gs[3, 0:2])

# 后续绘图逻辑不变
hist2d = ax1.hist2d(gaussian1, gaussian2, bins=100, density=True, cmap=plt.cm.jet)
fig.colorbar(hist2d[3], cax=ax4, orientation="horizontal")

内容的提问来源于stack exchange，提问作者Abel Gutiérrez