Python循环绘图问题：如何将4组数据绘制到独立子图

我有一个形状为(2,2,1000)的numpy数组，代表收入组、年龄组以及每组的1000个观测样本。我尝试用for循环绘制以下4种组合的图表：

1. 年龄≤18岁、收入i0；
2. 年龄≤18岁、收入i1；
3. 年龄≥18岁、收入i0；
4. 年龄≥18岁、收入i1，

期望得到4个并排的独立子图，但目前所有绘图都重叠在同一张图上。请问如何修改代码实现独立子图？

以下是我的代码：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

elasticity = np.random.rand(2,2,1000)
print(elasticity.shape)

income = ['i0','i1']
age_gr= ['<=18','>18']

for i in range(len(age_gr)):
    for j in range((len(income))):
        plt.plot(elasticity[:,j,:], elasticity[i,:,:])
        plt.subplot(i,j)
plt.show()

你的代码存在几个关键问题：

• plt.subplot()参数使用错误：它的标准格式是plt.subplot(nrows, ncols, index)，索引从1开始，你直接传循环的(i,j)（从0开始）会导致索引错误，且未指定布局行数/列数。
• 绘图顺序错误：应该先创建子图，再在对应子图上绘图，否则所有内容都会默认画在当前唯一的图上。
• 数组索引逻辑可能偏差：需要明确elasticity的维度对应关系，确保取到目标分组的样本。

修改后的代码如下：

import numpy as np
import matplotlib.pyplot as plt

elasticity = np.random.rand(2,2,1000)
print(elasticity.shape)

income = ['i0','i1']
age_gr= ['<=18','>18']

# 一次性创建2行2列的子图布局，获取子图对象数组
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 8))

for i in range(len(age_gr)):
    for j in range(len(income)):
        # 获取当前子图的对象
        ax = axes[i, j]
        # 根据维度对应关系取数据：这里假设elasticity维度为[年龄组, 收入组, 样本数]
        ax.plot(elasticity[i, j, :])
        # 设置子图标题
        ax.set_title(f"年龄{age_gr[i]}, 收入{income[j]}")
        # 设置坐标轴标签
        ax.set_xlabel("观测样本序号")
        ax.set_ylabel("弹性值")

# 自动调整子图间距，避免标题、标签重叠
plt.tight_layout()
plt.show()

关键修改说明： #

1. 使用plt.subplots()批量创建子图，通过axes数组精准控制每个子图，比循环调用plt.subplot()更直观可靠。
2. 先选中目标子图ax，再调用ax.plot()绘图，确保每个图表独立绘制在对应子图中。
3. 修正数组索引：如果你的elasticity维度是[收入组, 年龄组, 样本数]，则将elasticity[i,j,:]改为elasticity[j,i,:]即可匹配分组逻辑。
4. 添加标题、坐标轴标签和布局调整，提升图表可读性。

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