带有动态大小子图的代码重复：如何解决？

Question

我正在尝试保存由多个子图组成的图片：问题是我本来不知道我需要多少个子图（这取决于num_plots）所以我想获得一个在不同情况下生成不同图的函数：

如果 num_plots 等于或小于 3，则制作一排恰好 num_plots 个子图
如果超过 3 行，则生成所需行数（每行大小为 3）

问题是我无法在不避免代码重复的情况下做到这一点。这是我写的代码：

def foo(num_plots, ...):
    # ... obtain data to plot ('all_results')
    total_cols = 3 if num_plots > 2 else num_plots
    total_rows = math.ceil(num_plots / total_cols)
    fig, axs = plt.subplots(nrows=total_rows, ncols=total_cols, figsize=(7 * total_cols, 7 * total_rows),
                            constrained_layout=True)
    for i in range(num_plots):
        row = i // total_cols
        pos = i % total_cols

        # psi_name and psi_value depend on the parameters of the function
    
        if num_plots == 1:
            sns.kdeplot(all_results[:, i], color='magenta', ax=axs)
            axs.hist(all_results[:, i], bins=nbins)
            axs.axvline(x=psi_value, linestyle='--', color='red')
            axs.set_title(f"Estimation test for {psi_name} = {psi_value}")
        elif total_rows == 1:
            sns.kdeplot(all_results[:, i], color='magenta', ax=axs[pos])
            axs[pos].hist(all_results[:, i], bins=nbins)
            axs[pos].axvline(x=psi_value, linestyle='--', color='red')
            axs[pos].set_title(f"Estimation test for {psi_name} = {psi_value}")
        else:
            sns.kdeplot(all_results[:, i], color='magenta', ax=axs[row, pos])
            axs[row, pos].hist(all_results[:, i], bins=nbins)
            axs[row, pos].axvline(x=psi_value, linestyle='--', color='red')
            axs[row, pos].set_title(f"Estimation test for {psi_name} = {psi_value}")
    plt.savefig(token_hex(8))

如您所见，问题是根据需要的地块数量，我必须改用 axs 或 axs[pos] 或 axs[row, pos]。请注意，仅使用 axs[row, pos] 会给我一个错误。 如何解决这个代码重复问题？

以下是此函数生成的几个绘图示例：

如果 num_plots 为 1:

如果 num_plots 是 3:

如果 num_plots 是 4:

Answer 1

add_subplots() 顺序绘图技术是否适合您的目的？但是，它被设置为在单个图表上绘制多个图形，因此单个图形会根据图表的大小最大化。您将收到警告，但可以忽略。

import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure(figsize=(12,9))

x = np.linspace(-3, 3, 20)
cnt = 5
cols = 3
rows = round(cnt / cols,0)

for i in range(cnt):
    ax = fig.add_subplot(rows,cols,i+1)
    ax.plot(x, x**i)

Answer 2

如果您希望约束布局起作用，您可以创建一个 gridspec 并逐步添加到其中：

import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure(figsize=(12,9), constrained_layout=True)


x = np.linspace(-3, 3, 20)
cnt = 5
cols = 3
rows = round(cnt / cols,0)

gs = fig.add_gridspec(int(rows), int(cols))

for ind in range(cnt):
    i = int(np.floor(ind / cols))
    j = ind % cols

    ax = fig.add_subplot(gs[i, j])
    ax.plot(x, x**ind)

plt.show()

请注意，上面使用的未来 matplotlib 3.4，add_subplot 也将与 constrained_layout 一起使用，但尚未发布：https://github.com/matplotlib/matplotlib/pull/17347

带有动态大小子图的代码重复：如何解决？

Code duplication with subplots of dyinamic size: how to solve it?

python

matplotlib

subplot