如何在Matplotlib中为行数可变的子图统一设置参数？

问题场景 #

用户原代码如下：

import math
import matplotlib.pyplot as plt

dim=len(filtered_df[hu].unique())+len(filtered_df['Cat1'].unique())
plt.rcParams["figure.figsize"] = [10, 6]
plt.rcParams["figure.autolayout"] = True
ydim=math.ceil(dim/4)

fig, axis= plt.subplots(ydim, dim, squeeze=True)
fig.subplots_adjust(hspace=1, wspace=1)
for ax in axis:
    ax.set_ylim(0,100)
    ax.set_xlim(-1,4)

当ydim=1时，axis是形状为(4,)的一维数组，循环可直接遍历；但ydim=2时，axis变为形状(2,4)的二维数组，直接遍历会拿到每一行的子图数组，而非单个子图，导致报错。

解决方案 #

以下几种方法可以统一处理不同维度的axis对象：

方法1：使用数组的flat属性扁平化迭代

这是最简单的方式，不管axis是一维还是二维，flat都会返回一个一维迭代器，遍历所有子图：

for ax in axis.flat:
    ax.set_ylim(0,100)
    ax.set_xlim(-1,4)

如果存在子图数量不足（比如ydim*dim大于实际需要的子图数）的情况，flat会遍历到空的子图对象，此时可以加判断跳过：

for ax in axis.flat:
    if ax is not None:
        ax.set_ylim(0,100)
        ax.set_xlim(-1,4)

方法2：判断维度后分情况处理

如果不想依赖扁平化操作，可以先判断axis的维度，再选择遍历方式：

if axis.ndim == 1:
    for ax in axis:
        ax.set_ylim(0,100)
        ax.set_xlim(-1,4)
else:
    for row in axis:
        for ax in row:
            ax.set_ylim(0,100)
            ax.set_xlim(-1,4)

方法3：使用numpy.nditer迭代

借助numpy的迭代器遍历所有子图元素：

import numpy as np

for ax in np.nditer(axis, flags=['refs_ok']):
    ax.item().set_ylim(0,100)
    ax.item().set_xlim(-1,4)

这里flags=['refs_ok']用于处理对象数组，item()用来取出迭代器中的子图对象。

总结 #

优先推荐方法1，代码简洁且能覆盖所有场景，是处理Matplotlib子图遍历最常用的方式。

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