Matplotlib动态子图维度适配问题：避免索引错误

问题场景 #

当动态绘制N×M、单图、1×M或N×1布局的子图时，plt.subplots()的返回值axs维度会随行列数变化：

• 行列数都大于1时，返回二维数组
• 其中一行或一列数为1时，返回一维数组
  直接使用axs[N, M]二维索引会触发错误：

IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed

原代码示例：

fig, axs = plt.subplots(max_N, max_N)

for N in range(max_N):
    for M in range(max_N):
        axs[N, M].plot(...)

无分支判断的解决方案 #

方案1：强制二维数组（推荐）

调用plt.subplots()时添加squeeze=False参数，强制返回二维数组，无论行列数是否为1：

fig, axs = plt.subplots(max_N, max_N, squeeze=False)

for N in range(max_N):
    for M in range(max_N):
        axs[N, M].plot(...)

原理：默认squeeze=True会自动压缩长度为1的维度，设置为False后，始终保留二维结构，直接用二维索引即可。

方案2：用numpy统一维度

利用numpy.atleast_2d()将axs转换为至少二维的数组：

import numpy as np

fig, axs = plt.subplots(max_N, max_N)
axs = np.atleast_2d(axs)

for N in range(max_N):
    for M in range(max_N):
        axs[N, M].plot(...)

原理：一维数组会被转为(1, N)或(N, 1)的二维数组，二维数组则保持不变，兼容原有索引逻辑。

方案3：扁平化遍历子图

如果不需要区分行列位置，直接扁平化axs后遍历：

fig, axs = plt.subplots(max_N, max_N)

for ax in axs.flatten():
    ax.plot(...)

原理：flatten()将任意维度的axs转为一维数组，循环中直接操作每个子图对象，无需关注索引维度。

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