Python原子电荷计算函数的if条件语句是否可以简化？

你的代码完全可以简化，同时还能补上原代码的潜在漏洞、提升可维护性，优化后的代码如下：

首先是依赖导入和简化后的energy函数：

import math as m

def energy(q_ion, r_12, Q=-1, l=6):
    # 计算相互作用能
    return (Q * q_ion * m.erf(r_12 / l)) / r_12

然后是简化后的charge函数：

# 原子与对应离子电荷的映射表
ATOM_CHARGE_MAP = {
    "C": -4,
    "N": -3,
    "Cs": 1,
    "H": 1,
    "Ge": 2,
    "Pb": 2,
    "Si": 2,
    "Br": -1,
    "I": -1,
    "X": -2
}

def charge(atom, r_12):
    # 按原子类型返回相互作用能
    if atom not in ATOM_CHARGE_MAP:
        raise ValueError(f"不支持的原子类型：{atom}")
    return energy(ATOM_CHARGE_MAP[atom], r_12)

优化说明 #

• 原charge函数使用多个独立if判断，即便匹配到对应原子也会执行后续所有判断，存在不必要的性能损耗，改用字典映射后查找效率为O(1)，性能更优
• 原子与电荷的对应关系统一整理到字典中，后续新增/修改原子类型只需要修改字典即可，不需要新增if分支，可维护性大幅提升
• 补上了原代码的潜在漏洞：如果传入未定义的原子类型，原代码会因为变量E未定义抛出模糊的报错，现在会给出明确的错误提示
• energy函数去掉了不必要的中间变量，直接返回计算结果，逻辑更清晰；固定参数Q和l改为默认参数，后续调整参数时不需要修改函数内部逻辑
• 所有原有功能完全兼容，已有调用代码不需要做任何修改就能直接使用

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