2. 分子图数据获取与预处理

分子图神经网络的第一步并不是搭模型，而是把化学结构变成一张“机器能读的图”。这一步决定了后续所有特征、消息传递和预测的上限。本节围绕三件事展开：从常见分子文件格式中读取结构信息、把原子和键转成图节点与边、以及做必要的数据清洗与标准化。

一、三种主流格式的读取逻辑

SMILES 是最紧凑的线性表示，用字符串描述原子和连接关系，例如 CCO 表示乙醇，c1ccccc1 表示苯环。它的优点是体积小、易存储，缺点是隐式氢、芳香性、立体信息需要额外解析。SDF 是结构化数据文件，包含原子块（坐标、元素、电荷）和键块（连接、键级），还常带属性字段，适合三维构象和批量数据。MOL2 常见于分子对接和力场场景，除原子键外还带原子类型、部分电荷和子结构信息，适合需要电荷或力场参数的建模。

读取时建议统一走 RDKit 这类化学信息学工具，而不是自己写解析器。下面是一个从 SMILES 构建分子对象的示例：

from rdkit import Chem
from rdkit.Chem import AllChem

smi = "CC(=O)Oc1ccccc1C(=O)O"   # 阿司匹林
mol = Chem.MolFromSmiles(smi)
if mol is None:
    raise ValueError("SMILES 解析失败，需检查语法")

mol = Chem.AddHs(mol)            # 补显式氢
AllChem.EmbedMolecule(mol, randomSeed=42)
AllChem.MMFFOptimizeMolecule(mol)
print(mol.GetNumAtoms(), mol.GetNumBonds())

对 SDF 和 MOL2，可用 Chem.SDMolSupplier 与 Chem.MolFromMol2File 读取。注意 SDF 可能包含多条记录，需逐条判断是否为 None，并记录失败索引，便于后续排查。

二、从分子对象到分子图

分子图定义为 G = (V, E)：节点是原子，边是化学键。节点特征通常包括元素类型、原子序数、度、形式电荷、杂化方式、是否在环上、是否芳香、手性标签、氢原子数等；边特征包括键级（单/双/三/芳香）、是否共轭、是否在环上、立体构型等。把这些离散属性编码成向量，才能送入 GNN。

一个常见的做法是：元素用 one-hot 或可学习嵌入，数值特征做归一化，布尔特征用 0/1。下面给出节点与边特征的构建片段：

import torch
from rdkit.Chem import rdchem

def atom_features(atom):
    return [
        atom.GetAtomicNum(),
        atom.GetDegree(),
        atom.GetFormalCharge(),
        int(atom.GetIsAromatic()),
        int(atom.IsInRing()),
        atom.GetTotalNumHs(),
    ]

def bond_features(bond):
    bt = bond.GetBondType()
    return [
        int(bt == rdchem.BondType.SINGLE),
        int(bt == rdchem.BondType.DOUBLE),
        int(bt == rdchem.BondType.TRIPLE),
        int(bt == rdchem.BondType.AROMATIC),
        int(bond.GetIsConjugated()),
        int(bond.IsInRing()),
    ]

def mol_to_graph(mol):
    xs, edge_index, edge_attr = [], [], []
    for atom in mol.GetAtoms():
        xs.append(atom_features(atom))
    for bond in mol.GetBonds():
        i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        edge_index += [[i, j], [j, i]]
        edge_attr += [bond_features(bond), bond_features(bond)]
    return (torch.tensor(xs, dtype=torch.float),
            torch.tensor(edge_index, dtype=torch.long).t().contiguous(),
            torch.tensor(edge_attr, dtype=torch.float))

这里把每条键写成双向边，保证消息可以双向传递。若使用有向消息传递框架，也可只保留单向边并在聚合时对称处理。

三、数据清洗与标准化

真实数据往往不干净。常见问题包括：SMILES 语法错误、SDF 中价键异常、重复分子、盐和溶剂残留、芳香性表示不一致、电荷未中和、同位素或立体信息缺失。处理顺序建议如下：

• 解析过滤：丢弃 MolFromSmiles 返回 None 的记录，并记录失败原因。
• 去盐与去溶剂：用 rdMolStandardize 的 LargestFragmentChooser 保留最大片段，去掉 Na⁺、Cl⁻、水等。
• 电荷中和：用 Uncharger 中和可中和的电荷，避免同一分子因质子化状态不同被当成两个样本。
• 芳香性统一：用 Chem.SetAromaticity 统一感知模型，保证苯环在 SMILES 与 SDF 中编码一致。
• 去重：用规范 SMILES（Chem.MolToSmiles）作为唯一键，去除重复结构。
• 坐标检查：若下游需要三维构象，检查 SDF 是否含坐标，缺失时用 ETKDG 重新生成。

标准化之后，建议把每个分子转成统一的图对象并缓存为 .pt 或 .npz，避免每次训练重复解析。下面是一个批量处理的骨架：

from rdkit.Chem.MolStandardize import rdMolStandardize

def standardize(mol):
    mol = rdMolStandardize.LargestFragmentChooser().choose(mol)
    mol = rdMolStandardize.Uncharger().uncharge(mol)
    Chem.SetAromaticity(mol, Chem.AromaticityModel.AROMATICITY_RDKIT)
    return mol

seen = set()
graphs = []
for smi in raw_smiles:
    mol = Chem.MolFromSmiles(smi)
    if mol is None:
        continue
    mol = standardize(mol)
    key = Chem.MolToSmiles(mol)
    if key in seen:
        continue
    seen.add(key)
    graphs.append(mol_to_graph(mol))

四、一个可复用的检查清单