Python函数内批量保存不同命名的多张图表

我通过zip、annotate等函数生成12张独立图表，目前只能通过基于y轴标签（唯一标识符）编写多个if语句逐个保存，操作十分繁琐。希望在函数内为每张图表赋予唯一名称，并保存到同一文件夹中。当前绘图函数如下：

def PlotFromSelectedOutputs(df_3, pH_array, y_var, constants_dict_array):
    fig = plt.figure()
    ax = fig.add_axes([0,0,1,1])
    
    marker_dict = {'Ter':'d', 'Gos':'s', 'MS':'o'}
    color_dict = {'Ter':'b', 'Gos':'g', 'MS':'y'}
    markers = [marker_dict.get(key) for key in df_3.Lithology]
    colors = [color_dict.get(key) for key in df_3.Lithology]
    for _m, _c, _x, _y, _l in zip(markers, colors, df_3.pH, df_3[y_var], df_3.Lithology):
      ax.scatter(_x, _y, marker=_m, color=_c, linewidths = 1, label=_l)
      plt.xlabel("pH")
      plt.ylabel(y_var)
      
    legend_without_duplicate_labels(ax)
      
    plt.xlim([0, 14])
    plt.rcParams['figure.dpi'] = 300
    plt.rcParams.update({'figure.max_open_warning': 0})
    
    for constants_dict in constants_dict_array:
      _ = plt.plot(pH_array, constants_dict['sol'](np.array(pH_array)), 
                   constants_dict['style'], color=constants_dict['color'], label=constants_dict['label'])
      _ = plt.annotate(constants_dict['label'], xy=constants_dict['xy'], xytext=constants_dict['xytext'],
                       color=constants_dict['color'], arrowprops=dict(color=constants_dict['color'], arrowstyle='->'))

# Save figures to "Graphs" folder. ## THIS NEEDS TO BE IMPROVED ##
if y_var == "Log Activity of $Hg^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Hg2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Fe^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Fe2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Fe^{3+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Fe3p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Al^{3+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Al3p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Cu^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Cu2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Cd^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Cd2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Ni^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Ni2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Pb^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Pb2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Zn^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Zn2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $Mn^{2+}$ (mol/L)":
    plt.savefig('Graphs\\' + "Mn2p.png", bbox_inches ="tight")
if y_var == "Log Activity of $H_{2}AsO_{3}^{-}$ (mol/L)":
    plt.savefig('Graphs\\' + "As3p.png", bbox_inches ="tight")
if y_var == "Log Activity of $H_{2}AsO_{4}^{-}$ (mol/L)":
    plt.savefig('Graphs\\' + "As5p.png", bbox_inches ="tight")

优化方案 #

方法1：用映射字典替代多if判断 #

创建y轴标签与文件名的映射字典，直接通过键值对匹配获取文件名，代码更简洁易维护：

# 定义y_var到文件名的映射字典
y_var_filename_map = {
    "Log Activity of $Hg^{2+}$ (mol/L)": "Hg2p.png",
    "Log Activity of $Fe^{2+}$ (mol/L)": "Fe2p.png",
    "Log Activity of $Fe^{3+}$ (mol/L)": "Fe3p.png",
    "Log Activity of $Al^{3+}$ (mol/L)": "Al3p.png",
    "Log Activity of $Cu^{2+}$ (mol/L)": "Cu2p.png",
    "Log Activity of $Cd^{2+}$ (mol/L)": "Cd2p.png",
    "Log Activity of $Ni^{2+}$ (mol/L)": "Ni2p.png",
    "Log Activity of $Pb^{2+}$ (mol/L)": "Pb2p.png",
    "Log Activity of $Zn^{2+}$ (mol/L)": "Zn2p.png",
    "Log Activity of $Mn^{2+}$ (mol/L)": "Mn2p.png",
    "Log Activity of $H_{2}AsO_{3}^{-}$ (mol/L)": "As3p.png",
    "Log Activity of $H_{2}AsO_{4}^{-}$ (mol/L)": "As5p.png"
}

# 替换原有的所有if语句
filename = y_var_filename_map.get(y_var, f"unknown_{y_var.replace(' ', '_')}.png")
plt.savefig(f'Graphs\\{filename}', bbox_inches="tight")

说明：

• 字典清晰对应每个y轴标签和目标文件名，后续新增/修改只需调整字典，无需修改逻辑
• get方法提供默认值，避免因y_var未匹配导致报错，默认文件名自动将空格替换为下划线，保证文件命名合法性

方法2：自动从y_var提取信息生成文件名 #

如果y_var格式固定，可通过正则表达式自动提取离子名称和价态，无需手动维护映射字典：

import re

# 从y_var中提取离子核心信息
match = re.search(r'Log Activity of \$(.*?)\$', y_var)
if match:
    ion_str = match.group(1)
    # 匹配元素符号和价态
    elem_val_match = re.search(r'([A-Za-z]+|\w+\{.*?\})(\d+[+-]?)', ion_str)
    if elem_val_match:
        elem_raw = elem_val_match.group(1)
        valency = elem_val_match.group(2).replace('+', 'p').replace('-', 'm')
        # 特殊处理As的格式
        if 'AsO' in elem_raw:
            elem = 'As'
            valency = '3p' if '3' in ion_str else '5p'
        else:
            elem = elem_raw.strip('{}')
        filename = f"{elem}{valency}.png"
    else:
        filename = f"{ion_str.replace(' ', '_')}.png"
else:
    filename = f"unknown_{y_var.replace(' ', '_')}.png"

plt.savefig(f'Graphs\\{filename}', bbox_inches="tight")

说明：

• 通过正则自动解析y_var中的离子信息，无需手动维护映射关系
• 针对As的特殊格式做了适配，保证生成的文件名与原命名一致
• 包含异常处理，避免格式不匹配时出错

额外优化建议 #

• 提前检查Graphs文件夹是否存在，避免保存失败：

  import os
  os.makedirs('Graphs', exist_ok=True)  # 文件夹不存在则创建，存在则不报错
  

• 绘图完成后添加plt.close(fig)，释放内存，避免打开过多图表导致内存占用过高

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