Matplotlib子图尺寸、间距配置问题及优化诉求

核心需求 #

• 创建纵向排列的两个Matplotlib子图
• 子图A设置aspect='equal'
• 两个子图宽度保持一致
• 图形保存为pgf格式
• 设置图形最大尺寸

额外需求 #

• 可设置子图B的高度（绝对值或相对于子图A的相对值）
• 可设置子图A与B之间的间距（绝对值或相对于子图A的相对值）

当前遇到的问题 #

• 仅给子图A设置ax.set_aspect('equal')会导致A、B宽度不一致
• 使用plt.figure(figsize=(w,h))结合pgf后端时，图形会自动调整间隙填充给定尺寸，希望仅缩放适配尺寸而非拉伸
• 带负号的标签基线与无负号标签不一致的问题已确认属于pgf端问题

现有实现代码 #

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import sys

matplotlib.use('pgf')
matplotlib.rcParams.update({
    'text.usetex'     : True,
    'pgf.texsystem'   : 'pdflatex',
})

color_levels = 128
cmap = matplotlib.cm.RdBu_r
range_min = 0
range_max = np.pi

gkw = {'height_ratios':[0.03, 1] }
fig, axs = plt.subplots(2, 1, gridspec_kw = gkw)

lvlmax = 2
lvlmin = -2

ax = axs[1]
ax.set_xlim(range_min, range_max)
ax.set_ylim(range_min, range_max)
ax.set_aspect('equal')
ax.grid(False)

X = np.linspace(range_min, range_max, color_levels)
Z = np.linspace(range_min, range_max, color_levels)
X, Z = np.meshgrid(X, Z)
val = np.power(np.sin(X), 2) + np.cos(Z)

cf = ax.contourf(X, Z, val, cmap=cmap, zorder=1, extend='both',
                 levels=np.linspace(lvlmin, lvlmax, color_levels, endpoint = True))

axes_ratio = gkw['height_ratios'][1]/gkw['height_ratios'][0]

# colorbar
ax = axs[0]
ax.set_xlim(lvlmin, lvlmax)
y_max = (lvlmax-lvlmin)/30
ax.set_ylim(0, y_max)
ax.grid(False)

X = np.tile(np.linspace(lvlmin, lvlmax, color_levels), (10, 1))
Z = np.transpose(np.tile(np.linspace(0, y_max, 10), (color_levels, 1)))
color = np.tile(np.linspace(lvlmin, lvlmax, color_levels), (10, 1))

cf = ax.contourf(X, Z, color, cmap=cmap, zorder=1,
                    levels=np.linspace(lvlmin, lvlmax, color_levels, endpoint = True))

ax.tick_params(left=False, right=False , top=True, bottom=False,
                labelbottom=False, labeltop=True, labelleft=False)

y_axis_length = np.diff(np.array(ax.get_ylim()))[0]
x_axis_length = np.diff(np.array(ax.get_xlim()))[0]
other_axis_aspect = axs[1].get_aspect()

ax.set_aspect((other_axis_aspect/axes_ratio)/(y_axis_length/x_axis_length))

plt.subplots_adjust(wspace=0, hspace=0.05)

boxes = [ax.get_tightbbox() for ax in axs]
aspect = ((max([b.y1 for b in boxes]) - min([b.y0 for b in boxes])) /
          (max([b.x1 for b in boxes]) - min([b.x0 for b in boxes])))
fig.set_size_inches(3.4, 3.4*aspect)

fig.canvas.draw()
fig.savefig(sys.argv[1], bbox_inches = 'tight', transparent = True)

代码存在的问题 #

该代码可生成目标图形，但偶尔会出现子图间距异常的情况，推测是tightbbox未按预期工作，且整体实现方式较为生硬，希望找到更优雅的解决方案。

注：示例中用子图模拟颜色条，已知有专用颜色条命令，但自动颜色条存在尺寸问题，且等高线数量较多时颜色条会以png格式而非pgf格式保存，此外需求场景不限于颜色条。

内容的提问来源于stack exchange，提问作者fales