Matplotlib遍历df2参数在同一图中绘制多条线及函数返回值存储方法

核心修改点 #

• 补充自定义函数的返回值，将计算得到的wp、tempp、ycp、yap、pp五个数组作为返回值输出
• 修复原代码中随机数生成的逻辑错误，原来的循环没有将生成的随机数存入randomlist，会导致运算报错
• 把df2的初始化移到循环外，避免函数内部覆盖df2变量
• 遍历df2每一行拿到计算结果后，直接在对应子图上绘制线条，多次调用plot方法会自动在同一张子图上叠加多条线

import pandas as pd
import random
import matplotlib.pyplot as plt

def calculate_values(ConstantA, ConstantB, tst, temp, dtube):
    # 修复随机数生成逻辑，将生成的数存入列表
    randomlist = []
    for i in range(0, 5):
        n = random.randint(1, 30)
        randomlist.append(n)
    wp = ConstantA * pd.Series(randomlist)
    tempp = ConstantB * pd.Series(randomlist)
    ycp = tst * pd.Series(randomlist)
    yap = temp * pd.Series(randomlist)
    pp = dtube * pd.Series(randomlist)
    # 新增返回值，输出计算得到的五个数组
    return wp, tempp, ycp, yap, pp

# 把df2初始化移到循环外，这里示例构造测试数据，你可以替换成自己的真实数据生成逻辑
df2 = pd.DataFrame([
    [1, 2, 3, 4, 5],
    [2, 3, 4, 5, 6],
    [3, 4, 5, 6, 7],
    [4, 5, 6, 7, 8],
    [5, 6, 7, 8, 9]
], columns=['ConstantA', 'ConstantB', 'tst', 'temp', 'dtube'])

# 提前创建子图对象
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10,8))
fig.suptitle('Cooled;Tst=523')

# 遍历每一行计算并直接绘图
for index, row in df2.iterrows():
    wp, tempp, ycp, yap, pp = calculate_values(row['ConstantA'], row['ConstantB'], row['tst'], row['temp'], row['dtube'])
    # 每次遍历都在对应子图新增一条线
    ax1.plot(wp, tempp, label=f'组{index+1}')
    ax2.plot(wp, ycp, label=f'组{index+1}')
    ax3.plot(wp, yap, label=f'组{index+1}')
    ax4.plot(wp, pp, label=f'组{index+1}')

# 设置坐标轴标签
ax1.set_ylabel('T(K)')
ax1.legend()
ax2.set_ylabel('y(C%)')
ax2.legend()
ax3.set_ylabel('yA ()')
ax3.set_xlabel('w (1000 ka)')
ax3.legend()
ax4.set_ylabel('E (bar)')
ax4.set_xlabel('w (1000 kg)')
ax4.legend()

plt.tight_layout()
plt.subplots_adjust(top=0.9)
plt.show()

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