如何用RDKit绘制分子极性表面积？解决芳香环误高亮问题

问题根源 #

你当前的代码用Gasteiger负电荷原子标记PSA区域，但Gasteiger电荷是经验性的部分电荷计算，芳香环碳原子因共轭效应可能被算出负电荷，而这些原子并不属于PSA的贡献部分（PSA仅关注极性杂原子及其相连氢），因此出现错误高亮。

正确思路：基于PSA定义识别原子 #

极性表面积（PSA）的计算仅考虑分子中的极性杂原子（O、N、P、S）以及这些原子上连接的氢原子，我们可以直接基于原子类型和连接关系筛选这些原子，而非依赖电荷计算。

改进后的代码 #

from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem import AllChem

def highlight_psa_atoms(mol):
    psa_atoms = []
    for atom in mol.GetAtoms():
        atom_num = atom.GetAtomicNum()
        # 识别PSA相关杂原子：O, N, P, S
        if atom_num in [7, 8, 15, 16]:
            psa_atoms.append(atom.GetIdx())
            # 添加这些原子上连接的氢原子
            for nbr in atom.GetNeighbors():
                if nbr.GetAtomicNum() == 1:
                    psa_atoms.append(nbr.GetIdx())
    # 去重避免重复添加同一个氢
    psa_atoms = list(set(psa_atoms))
    # 设置红色高亮样式
    highlight_style = {atom_id: (1, 0, 0) for atom_id in psa_atoms}
    return highlight_style

# 示例分子
smiles = "CC(=O)OC1=CC=CC=C1C(O)=O"
mol = Chem.MolFromSmiles(smiles)

# 获取高亮样式
highlight_style = highlight_psa_atoms(mol)

# 绘制分子
img = Draw.MolToImage(mol, size=(300, 300), highlightAtoms=highlight_style, wedgeBonds=True, kekulize=True, wedgeLineWidth=2)
img

额外优化：匹配RDKit官方TPSA计算逻辑 #

如果需要和RDKit的Descriptors.TPSA结果完全对齐，可以复用RDKit内部的原子贡献规则，确保高亮原子与官方计算的PSA贡献原子完全一致：

from rdkit.Chem import Descriptors, rdMolDescriptors

def highlight_psa_atoms_tpsa_aligned(mol):
    # 获取每个原子的TPSA贡献值
    atom_contribs = rdMolDescriptors._CalcTPSAContribs(mol)
    # 筛选贡献值大于0的原子（即参与PSA的原子）
    psa_atoms = [idx for idx, contrib in enumerate(atom_contribs) if contrib > 0]
    highlight_style = {atom_id: (1, 0, 0) for atom_id in psa_atoms}
    return highlight_style

# 使用示例
highlight_style = highlight_psa_atoms_tpsa_aligned(mol)
img = Draw.MolToImage(mol, size=(300, 300), highlightAtoms=highlight_style, wedgeBonds=True, kekulize=True, wedgeLineWidth=2)
img

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