如何让matplotlib.pyplot.subplots返回的ax始终支持下标索引

调用plt.subplots()时指定squeeze=False即可，这个参数会强制返回的ax对象保持二维数组结构，无论子图行数、列数是多少，都可以统一用ax[i,j]的下标形式访问。

代码示例 #

import matplotlib.pyplot as plt
# 不管nrows和ncols设为1还是更大的数值，ax始终是二维数组
fig, ax = plt.subplots(nrows=1, ncols=1, squeeze=False)
# 原有的索引写法完全可用，不会报错
ax[0,0].plot([1,2,3],[2,4,8])
plt.show()

原理解释 #

默认情况下squeeze=True，matplotlib会自动压缩多余的数组维度：

• 1x1布局时直接返回单个Axes对象，不支持下标索引
• 单行/单列布局时返回一维数组，只能用ax[i]的形式索引
• 行列数均大于1时才返回二维数组，支持ax[i,j]索引

设置squeeze=False后，所有场景返回的ax都是维度为(nrows, ncols)的二维数组，完全统一索引写法，不需要针对不同布局调整代码。

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