Можете ли вы заставить wspace и hspace подзаголовков фигуры иметь фиксированное значение в matplotlib независимо от размера фигуры

Я пытаюсь создать функцию для построения нескольких изображений в сетке с одной цветовой полосой и гистограммой. Я хотел бы, чтобы расстояние между всеми графиками было фиксированным значением, а цветная полоса охватывала высоту всех изображений, а гистограмма - ширину изображений/цветной полосы. У меня есть некоторый код, который работает, но для его работы требуется, чтобы размер фигуры был установлен на определенное соотношение сторон. Это не идеально, потому что я хочу использовать эту функцию для изображений с разным соотношением сторон и для разного количества изображений 2x1, 1x2, 2x2 и т. д.

Этот код выводит 3 цифры с разным соотношением сторон. Я хотел бы, чтобы какой-либо избыточный размер применялся к интервалу между границами, а не к подзаголовку wspace, hspace spacing.

ширина риса: https://i.stack.imgur.com/BB1Cz.png

высота инжира: https://i.stack.imgur.com/G5C34.png

рис хороший: https://i.stack.imgur.com/AVX6C.png

Вот код:

import math
import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt


def compare_frames(frames, columns, bins=256, alpha=.5, vmin=None, vmax=None, fig=None):
    if vmin is None:
        vmin = min([f.min() for f in frames])
    if vmax is None:
        vmax = max([f.max() for f in frames])
    if fig == None:
        fig = plt.figure()
    color_cycle = plt.get_cmap('tab10')
    rows = math.ceil(len(frames)/columns)
    width_ratios = [1 for col in range(columns)] + [.05]
    
    gs = mpl.gridspec.GridSpec(rows + 1, columns + 1, figure=fig, width_ratios=width_ratios)
    images = []
    for row in range(rows):
        for col in range(columns):
            idx = row*columns + col
            if idx < len(frames):
                ax = fig.add_subplot(gs[row, col])
                ax.get_xaxis().set_ticks([])
                ax.get_yaxis().set_ticks([])
                for spine in ['bottom', 'top', 'left', 'right']:
                    ax.spines[spine].set_color(color_cycle(idx))
                    ax.spines[spine].set_linewidth(3)
                images.append(ax.imshow(frames[idx], vmin=vmin, vmax=vmax))
    
    cax = fig.add_subplot(gs[0:-1, -1])
    plt.colorbar(images[0], cax=cax)

    hax = fig.add_subplot(gs[-1, :])
    for i, frame in enumerate(frames):
        hax.hist(frame.ravel(), bins=256, range=(vmin, vmax), color=color_cycle(i), alpha=alpha)
    
    fig.subplots_adjust(wspace=.05, hspace=.05)


if __name__ == '__main__':
    x_size = 640
    y_size = 512
    frames = []
    for i in range(4):
        frames.append(np.random.normal(i + 1, np.sqrt(i + 1), size=(y_size, x_size)))
    
    fig_wide = plt.figure(figsize=(12, 8))
    compare_frames(frames, 2, fig=fig_wide)
    fig_tall = plt.figure(figsize=(6, 8))
    compare_frames(frames, 2, fig=fig_tall)
    fig_nice = plt.figure(figsize=(6.9, 8))
    compare_frames(frames, 2, fig=fig_nice)
    plt.show()