如何将垂直堆叠子图排列为3×2网格？

问题描述 #

现有6个由stacked_lineplots()生成的垂直堆叠子图（每个子图包含2个上下排列的折线图），需要将它们排列为2行3列的网格布局，且必须通过独立函数arrange_stacked_subfigures_into_3x2_grid(figures)完成。

完整解决方案代码 #

import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

def arrange_stacked_subfigures_into_3x2_grid(figures):
    # 创建新的总画布，尺寸适配3列2行的布局
    grid_fig = plt.figure(figsize=(24, 12))
    # 定义2行3列的主网格
    main_gs = grid_fig.add_gridspec(2, 3)
    
    for idx, sub_fig in enumerate(figures):
        # 计算当前子图在主网格中的行、列位置
        row = idx // 3
        col = idx % 3
        # 在主网格单元格内创建2行1列的子网格，对应原堆叠子图的两个垂直轴
        sub_gs = main_gs[row, col].subgridspec(2, 1)
        
        # 获取原堆叠子图中的两个坐标轴
        original_axes = sub_fig.get_axes()
        for ax_idx, original_ax in enumerate(original_axes):
            # 在子网格中创建新坐标轴
            new_ax = grid_fig.add_subplot(sub_gs[ax_idx])
            
            # 复制原坐标轴的折线数据
            for line in original_ax.get_lines():
                new_ax.plot(line.get_xdata(), line.get_ydata(), 
                           color=line.get_color(), label=line.get_label())
            
            # 复制标题、坐标轴标签
            new_ax.set_title(original_ax.get_title())
            new_ax.set_xlabel(original_ax.get_xlabel())
            new_ax.set_ylabel(original_ax.get_ylabel())
            
            # 复制图例（如果存在）
            if original_ax.get_legend():
                new_ax.legend()
            
            # 复制坐标轴刻度
            new_ax.set_xticks(original_ax.get_xticks())
            new_ax.set_yticks(original_ax.get_yticks())
    
    # 调整整体布局，避免元素重叠
    grid_fig.tight_layout()
    # 关闭原画布释放内存
    for fig in figures:
        plt.close(fig)
    return grid_fig

def stacked_lineplots():
    # 创建包含2个垂直堆叠子图的画布
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6))
    
    # 模拟数据
    x = [0, 1, 2, 3, 4]
    y1 = [0, 1, 4, 9, 16]
    y2 = [16, 9, 4, 1, 0]
    
    # 绘制折线
    ax1.plot(x, y1, label='y = x^2')
    ax2.plot(x, y2, label='y = 16 - x^2', color='orange')
    
    # 调整子图布局
    plt.tight_layout()
    
    return fig

# 生成6个堆叠子图
six_figures = [stacked_lineplots() for _ in range(6)]

# 排列为3×2网格
grid_figure = arrange_stacked_subfigures_into_3x2_grid(six_figures)

# 可选：保存最终网格图
# grid_figure.savefig('3x2_stacked_grid.png', dpi=100, bbox_inches='tight')

print("Done")

关键实现说明 #

• 网格布局设计：通过GridSpec先创建2行3列的主网格，每个主单元格再嵌套2行1列的子网格，完美适配原堆叠子图的垂直结构。
• 内容复制：由于Matplotlib无法直接移动坐标轴，通过复制原坐标轴的折线数据、标签、图例、刻度等属性，确保新画布与原堆叠子图内容一致。
• 资源优化：处理完成后关闭原画布，避免内存占用过高。

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