Matplotlib使用plt.subplots时如何优雅处理可变数量坐标轴

问题根因 #

调用plt.subplots()创建子图网格时，默认开启squeeze=True参数，会自动压缩返回的坐标轴数组维度，不同子图数量下返回值结构不一致：

• 当p=1、q=1（单子图）时，直接返回单个Axes对象，不支持数组索引
• 当p=1或q=1（单行/单列子图）时，返回一维Axes数组，仅支持单参数索引
• 仅当p>1且q>1时，才返回p行q列的二维Axes数组，支持[行索引, 列索引]的二维访问方式
  这就是直接使用ax[ii, jj]写法在单图、单行/单列场景触发类型错误的核心原因。

最优解决方案 #

创建子图时传入squeeze=False参数，即可强制让返回的ax永远保持p行q列的二维数组结构，无需编写任何分支判断逻辑，全程统一使用ax[ii, jj]的二维索引写法访问所有子图坐标轴即可。

修正后可直接运行的代码 #

from matplotlib import pyplot as plt
import numpy as np
from random import randrange

N = 100
x = np.arange(0,100,1)
y = np.random.rand(len(x),1)

p = randrange(4)+1
q = randrange(4)+1

# 核心改动：添加squeeze=False，固定返回二维坐标轴数组
fig, ax = plt.subplots(p, q, sharex=True, squeeze=False)
for ii in range(p):
    for jj in range(q):
        ax[ii, jj].plot(x, y)
        ax[ii, jj].grid(True)
        ax[ii, jj].set_xlim([-3, 5])
plt.suptitle('Test')
plt.show()

更简洁的遍历写法 #

如果不需要按行列坐标定位特定子图，只是要对所有子图执行相同的绘图/属性设置操作，在设置squeeze=False的前提下，可以直接用ax.flat做扁平化遍历，省去双层循环：

from matplotlib import pyplot as plt
import numpy as np
from random import randrange

N = 100
x = np.arange(0,100,1)
y = np.random.rand(len(x),1)

p = randrange(4)+1
q = randrange(4)+1

fig, ax = plt.subplots(p, q, sharex=True, squeeze=False)
# 扁平化遍历所有子图，自动适配任意p、q取值
for ax_item in ax.flat:
    ax_item.plot(x, y)
    ax_item.grid(True)
    ax_item.set_xlim([-3, 5])
plt.suptitle('Test')
plt.show()

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