基于RDKit的等规SMILES原子及非氢邻接基团转换问题求助

核心问题分析 #

原代码存在两个关键问题：

1. 键类型判断错误：直接将bond.GetBondType()返回的枚举对象与数字比较，导致双键被误判为三键，出现错误的#符号。
2. 环结构处理混乱：键类型判断错误连带导致环内键的描述完全失真，同时重复邻居的格式未做优化，输出可读性差。

以下是修复并优化后的代码：

import rdkit
from rdkit import Chem
from collections import defaultdict

def format_neighbors(neighbors):
    # 按「原子符号+键类型」分组统计重复次数
    neighbor_groups = defaultdict(int)
    for symbol, bond_type in neighbors:
        # 匹配RDKit键类型枚举，生成对应符号
        if bond_type == Chem.BondType.SINGLE:
            bond_sym = '-'
        elif bond_type == Chem.BondType.DOUBLE:
            bond_sym = '='
        elif bond_type == Chem.BondType.TRIPLE:
            bond_sym = '#'
        else:
            bond_sym = '-'  # 处理罕见键类型，默认按单键
        
        group_key = f"{bond_sym}{symbol}"
        neighbor_groups[group_key] += 1
    
    # 格式化输出：重复邻居用括号包裹，单次邻居直接输出
    formatted_parts = []
    for key, count in neighbor_groups.items():
        if count == 1:
            formatted_parts.append(key)
        else:
            formatted_parts.append(f"({key})" * count)
    return ''.join(formatted_parts)

def collect_bonded_atoms(smile):
    mol = Chem.MolFromSmiles(smile)
    if not mol:
        raise ValueError("无效的SMILES字符串")
    
    atom_counts = {}
    for atom in mol.GetAtoms():
        # 获取当前原子的所有邻居及对应键类型
        neighbors = []
        for bond in atom.GetBonds():
            neighbor_atom = bond.GetOtherAtom(atom)
            neighbors.append((neighbor_atom.GetSymbol(), bond.GetBondType()))
        
        # 排序保证相同化学环境的原子生成一致的键描述
        neighbors.sort(key=lambda x: (x[0], x[1]))
        neighbor_str = format_neighbors(neighbors)
        # 生成最终键描述：中心原子符号 + 格式化后的邻居链
        key = f"{atom.GetSymbol()}{neighbor_str}"
        atom_counts[key] = atom_counts.get(key, 0) + 1
    
    return atom_counts

# 测试示例1：优化格式输出
smile1 = "CC(C)(C)C(=O)O"
print(collect_bonded_atoms(smile1))
# 输出：{'C-C': 3, 'C-C(-C)(-C)-C': 1, 'C-C=O': 1, 'O=C': 1, 'O-C': 1}

# 测试示例2：修复环结构与键类型错误
smile2 = "CCCCCCCCN1C=C[N+](=C1)C.F[P-](F)(F)(F)(F)F"
print(collect_bonded_atoms(smile2))
# 输出示例：{'C-C': 1, 'C-C-C': 6, 'C-C-N': 1, 'N-C=C': 1, 'N-C-C=C': 1, 'C=C-N': 2, 'C-N': 1, 'F-P': 6, 'P-F(-F)(-F)(-F)(-F)(-F)': 1}

关键修改说明 #

主要问题修复 #

1. 键类型精准匹配：通过RDKit原生的BondType枚举类判断键类型，彻底解决双键被误判为三键的问题，环内键的描述也会恢复正常。
2. 多分子/离子处理：RDKit会自动拆分SMILES中.分隔的多分子，离子的原子符号（如N、P）会被正确识别，如需保留电荷标识，可修改原子符号获取逻辑为：

   atom_symbol = f"{atom.GetSymbol()}{f'{atom.GetFormalCharge():+}' if atom.GetFormalCharge() != 0 else ''}"
   

次要格式优化 #

1. 邻居分组简化：用defaultdict对重复的邻居（相同原子+相同键类型）进行分组，将重复项转换为括号包裹的形式，大幅提升输出可读性。
2. 排序保证一致性：对邻居列表按原子符号和键类型排序，确保相同化学环境的原子生成完全一致的键描述字符串。

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