在 for 循环中与 Matplotlib 子图作斗争

Question

我有以下代码：

# Make plot of channels with gaps
fig, ax = plt.subplots(nrows=len(gap_list), ncols=1, figsize=(12,len(gap_list)), sharex=True, squeeze=False)

for ch in gap_list:
    i = gap_list.index(ch)
    resample_s = 4*ch_gap[ch]['rec_rate']
    ylabel = ch + ' (' + ch_gap[ch]['board'] +') - '+ ch_gap[ch]['unit']
    data = df[ch].resample(f'{resample_s}s').mean()
    is_nan = data.isnull()
    ax[i].fill_between(data.index, 0, (is_nan*data.max()), color='r', step='mid', linewidth='0')
    ax[i].plot(data.index, data, color='b', linestyle='-', marker=',', label=ylabel)
    ax[i].legend(loc='upper left')


plt.tight_layout()
plt.show()

其中 gap_list 是一个列表，其中包含来自 pandas 数据框 (df) 的列名。列表的长度可以是 1 到 10 之间的任何值。当 nrows > 1 时它可以正常工作。但是当 nrows == 1 时，我会遇到一个引发异常的问题：

'AxesSubplot' object is not subscriptable

然后我找到了squeeze kwarg并将其设置为false，我认为一切都很好，但现在代码引发了这个异常：

'numpy.ndarray' object has no attribute 'fill_between'

然后我采取了不同的策略并将图形设置在循环之外并将子图创建放在循环内：

fig = plt.figure(figsize=(12,len(gap_list)))

在for循环中创建的每个轴如下：

ax = plt.subplot(len(gap_list), 1, i+1)

这适用于 nrows=1 和 norws > 1。但是，我找不到使所有子图共享 X 轴的好方法。在原始方法中，我可以为 plt.subplots() 设置 sharex=True。

所以感觉原始方法更符合实际，但缺少一个要素来更好地处理 nrows=1 案例。

Answer 1

我认为保留原始代码最简单，但只需检查 ax 是否为 numpy 数组。

当nrows > 1 时，ax 将是 matplotlib 轴的 numpy 数组，因此索引到 ax。当nrows == 1，ax 将只是matplotlib 的轴，所以直接使用它。

import numpy as np

...

for ch in gap_list:
    
    ...
    
    # if `ax` is a numpy array then index it, else just use `ax`
    ax_i = ax[i] if isinstance(ax, np.ndarray) else ax

    # now just use the `ax_i` handle
    ax_i.fill_between(data.index, 0, (is_nan*data.max()), color='r', step='mid', linewidth='0')
    ax_i.plot(data.index, data, color='b', linestyle='-', marker=',', label=ylabel)
    ax_i.legend(loc='upper left')