如何在Python中基于条件循环绘制多图并事后保存

需求说明 #

在for循环内根据不同条件将多组数据绘制到不同图表，循环结束后为各图表添加坐标轴标签、图例等标注，再保存为文件。以下是对应MATLAB实现代码，提供Python matplotlib的解决方案。

MATLAB参考代码 #

A = [some 2d matrix]
B = [some 2d matrix]

[row, col] = size(A)
for ii = 1:row
    x = A(ii,:)
    y = B(ii,:)
    
    if (some condition)
        figure(1)
        plot(x,y)
        hold on
    elseif (some condition)
        figure(2)
        plot(x,y)
        hold on
    elseif (some condition)
        figure(3)
        plot(x,y)
        hold on
    else
        figure(4)
        plot(x,y)
        hold on
    end
end

figure(1)
xlabel(...)
ylabel(...)
legend(...)
saveas()

figure(2)
xlabel(...)
ylabel(...)
legend(...)
saveas()

figure(3)
xlabel(...)
ylabel(...)
legend(...)
saveas()

figure(4)
xlabel(...)
ylabel(...)
legend(...)
saveas()

Python matplotlib解决方案 #

核心思路是提前创建所有需要的画布(Figure)和坐标轴(Axes)对象，循环内直接调用对应Axes绘图，最后统一配置标注并保存。

import matplotlib.pyplot as plt
import numpy as np

# 模拟测试数据，替换为你的真实数据
A = np.random.rand(20, 10)  # 20行10列的随机矩阵
B = np.random.rand(20, 10)

# 提前创建4个Figure和对应的Axes，对应MATLAB的figure(1)到figure(4)
fig1, ax1 = plt.subplots()
fig2, ax2 = plt.subplots()
fig3, ax3 = plt.subplots()
fig4, ax4 = plt.subplots()

# 存储每条曲线的标签，用于后续生成图例（按需调整）
labels1 = []
labels2 = []
labels3 = []
labels4 = []

row, col = A.shape
for ii in range(row):
    x = A[ii, :]
    y = B[ii, :]
    
    # 替换为你的真实条件判断逻辑
    if ii % 4 == 0:
        # 绘制到第一个图表
        ax1.plot(x, y)
        labels1.append(f'Curve {ii+1}')
    elif ii % 4 == 1:
        # 绘制到第二个图表
        ax2.plot(x, y)
        labels2.append(f'Curve {ii+1}')
    elif ii % 4 == 2:
        # 绘制到第三个图表
        ax3.plot(x, y)
        labels3.append(f'Curve {ii+1}')
    else:
        # 绘制到第四个图表
        ax4.plot(x, y)
        labels4.append(f'Curve {ii+1}')

# 统一配置各图表的标注并保存
# 图表1
ax1.set_xlabel('X轴标签')
ax1.set_ylabel('Y轴标签')
ax1.legend(labels1)
fig1.savefig('figure1.png', dpi=300, bbox_inches='tight')
plt.close(fig1)  # 释放资源

# 图表2
ax2.set_xlabel('X轴标签')
ax2.set_ylabel('Y轴标签')
ax2.legend(labels2)
fig2.savefig('figure2.png', dpi=300, bbox_inches='tight')
plt.close(fig2)

# 图表3
ax3.set_xlabel('X轴标签')
ax3.set_ylabel('Y轴标签')
ax3.legend(labels3)
fig3.savefig('figure3.png', dpi=300, bbox_inches='tight')
plt.close(fig3)

# 图表4
ax4.set_xlabel('X轴标签')
ax4.set_ylabel('Y轴标签')
ax4.legend(labels4)
fig4.savefig('figure4.png', dpi=300, bbox_inches='tight')
plt.close(fig4)

关键说明 #

• 无需使用hold on：matplotlib的Axes对象会自动保留所有绘制的曲线，每次调用plot()都会在同一坐标轴上新增曲线。
• 提前创建Figure/Axes：避免循环内重复创建，确保所有绘图都指向正确的图表。
• 统一配置标注：循环结束后集中设置标签、图例，逻辑更清晰，也便于统一调整格式。
• 保存后关闭Figure：使用plt.close()释放内存，避免资源占用。

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