3. 分子图构建：从SMILES到PyG的Data对象
这一章，咱们来动手把分子变成图。
说白了，就是把化学式、结构式，转成计算机能理解的数据结构。
我刚开始做这个的时候，总觉得这步很简单。结果一上手，坑不少。比如原子特征选多少？键要不要加方向？图太大显存不够怎么办？
别急，一个个聊。
3.1 RDKit解析分子：SMILES -> RDKit Mol
RDKit是个化学信息学工具包，功能极强。我个人习惯用它来处理分子。
先看最简单的操作：从SMILES字符串生成一个分子对象。
from rdkit import Chem
smiles = "CC(=O)OC1=CC=CC=C1C(=O)O" # 阿司匹林
mol = Chem.MolFromSmiles(smiles)
print(mol.GetNumAtoms()) # 输出：13
print(mol.GetNumBonds()) # 输出：13嗯，就这么简单。
但注意，RDKit默认会保留氢原子吗？不会。
上面的代码，mol里不含氢。如果你想显式加氢，得这么写：
mol = Chem.AddHs(mol)
print(mol.GetNumAtoms()) # 输出：21为什么我提这个？
因为很多初学者直接拿带H的分子去做图神经网络，特征维度暴增，训练效果反而不好。我建议，除非你的任务需要氢原子信息（比如某些力场计算），否则就不要加氢。
Chem.MolFromSmiles 得到分子，不加氢；然后用 Chem.RemoveHs(mol) 再确认一下。稳。
3.2 原子与键特征：选什么？怎么选？
特征选择，是分子图构建的核心。
图神经网络本质是在节点和边上做消息传递。节点特征决定了模型能“看到”什么信息。边特征决定了模型能“感知”什么连接关系。
常用原子特征（节点特征）：
| 特征 | 描述 | 维度 |
|---|---|---|
| 原子类型 | 如C、N、O、S等，one-hot编码 | 常见元素数（如10） |
| 原子度数 | 连接的键数（不含H） | 1维整数 |
| 形式电荷 | 原子的形式电荷 | 1维整数 |
| 杂化方式 | sp、sp2、sp3等 | one-hot编码 |
| 在环中 | 是否在环上 | 1维布尔值 |
| 手性 | 手性类型 | one-hot编码 |
| 氢原子数 | 连接的氢原子数 | 1维整数 |
常用键特征（边特征）：
| 特征 | 描述 | 维度 |
|---|---|---|
| 键类型 | 单键、双键、三键、芳香键 | one-hot编码（4维） |
| 共轭 | 是否共轭键 | 1维布尔值 |
| 在环中 | 是否在环上 | 1维布尔值 |
| 立体化学 | 顺反、E/Z等 | one-hot编码 |
我在项目里遇到过一个问题：特征维度选太多，模型训练慢，反而过拟合。
其实，特征不是越多越好。你想想看，一个分子就几十个原子，你给每个原子搞个上百维的特征，稀稀拉拉的，模型学不到东西。
这里给出一段代码，生成原子和键特征：
from rdkit import Chem
from rdkit.Chem import AllChem
def atom_features(atom):
# 原子类型 one-hot
atom_type = atom.GetAtomicNum()
atom_type_onehot = [0]*10
if atom_type in [6,7,8,9,16,17,35,53]:
idx = [6,7,8,9,16,17,35,53].index(atom_type)
atom_type_onehot[idx] = 1
# 度数
degree = atom.GetDegree()
# 形式电荷
formal_charge = atom.GetFormalCharge()
# 杂化方式
hybridization = atom.GetHybridization()
hyb_onehot = [0]*5
if hybridization in [Chem.rdchem.HybridizationType.SP,
Chem.rdchem.HybridizationType.SP2,
Chem.rdchem.HybridizationType.SP3]:
idx = [Chem.rdchem.HybridizationType.SP,
Chem.rdchem.HybridizationType.SP2,
Chem.rdchem.HybridizationType.SP3].index(hybridization)
hyb_onehot[idx] = 1
# 在环中
in_ring = 1 if atom.IsInRing() else 0
return atom_type_onehot + [degree, formal_charge] + hyb_onehot + [in_ring]
def bond_features(bond):
# 键类型
bond_type = bond.GetBondType()
btype_onehot = [0]*4
if bond_type in [Chem.rdchem.BondType.SINGLE,
Chem.rdchem.BondType.DOUBLE,
Chem.rdchem.BondType.TRIPLE,
Chem.rdchem.BondType.AROMATIC]:
idx = [Chem.rdchem.BondType.SINGLE,
Chem.rdchem.BondType.DOUBLE,
Chem.rdchem.BondType.TRIPLE,
Chem.rdchem.BondType.AROMATIC].index(bond_type)
btype_onehot[idx] = 1
# 共轭
is_conjugated = 1 if bond.GetIsConjugated() else 0
# 在环中
in_ring = 1 if bond.IsInRing() else 0
return btype_onehot + [is_conjugated, in_ring]3.3 图数据结构：PyG的Data对象
RDKit处理完了分子，特征也提取好了。接下来，怎么把这些数据喂给图神经网络？
我的选择是PyG（PyTorch Geometric）。PyG的Data对象是图数据的标准容器。
一个Data对象需要这么几个字段：
x: 节点特征矩阵，形状为 [N, F]edge_index: 边索引，形状为 [2, E]，注意是无向图要存双向edge_attr: 边特征矩阵，形状为 [E, D]y: 标签（如果有的话）
举个例子：一个分子有10个原子，13条键。节点特征F=10维，边特征D=7维。
那么：
x的形状是 [10, 10]edge_index的形状是 [2, 26]（双向就是13*2=26条边）edge_attr的形状是 [26, 7]
构建过程我直接贴代码：
import torch
from torch_geometric.data import Data
def mol_to_pyg_data(smiles, label=None):
mol = Chem.MolFromSmiles(smiles)
mol = Chem.RemoveHs(mol)
# 节点特征
node_features = []
for atom in mol.GetAtoms():
node_features.append(atom_features(atom))
x = torch.tensor(node_features, dtype=torch.float)
# 边索引和边特征
edge_index = []
edge_attr = []
for bond in mol.GetBonds():
i = bond.GetBeginAtomIdx()
j = bond.GetEndAtomIdx()
# 无向图，两个方向都要存
edge_index.append([i, j])
edge_index.append([j, i])
edge_attr.append(bond_features(bond))
edge_attr.append(bond_features(bond))
edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
edge_attr = torch.tensor(edge_attr, dtype=torch.float)
data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
if label is not None:
data.y = torch.tensor([label], dtype=torch.float)
return data
# 测试
smiles = "CC(=O)OC1=CC=CC=C1C(=O)O" # 阿司匹林
data = mol_to_pyg_data(smiles)
print(data)输出类似这样：
Data(x=[13, 10], edge_index=[2, 26], edge_attr=[26, 7])到这里，一个分子就被封装成一个图了。
3.4 实战中的几个问题
这部分内容不算理论，全是我踩过的坑。分享一下：
问题一：分子尺寸不一
分子大小不同，图的大小也不同。PyG的Data对象天然支持变尺寸，这一点很方便。但训练时要用DataLoader，它会自动处理批次。
问题二：坐标信息要不要加？
有些任务需要分子的3D构象，比如预测蛋白-配体结合能。这时候可以把原子坐标也作为特征加进去。但注意，构象可能有多个，要选哪个？
我个人建议：如果你的数据里有3D构象（比如来自晶体结构），可以用；如果只是从SMILES生成，那构建的构象可能不准，反而引入噪声。
问题三：数据泄露
这绝对是大坑。特别是做ADMET预测时，同一个化合物的类似物可能同时出现在训练集和测试集。如果你不处理，模型看似精度很高，但实际泛化能力很差。
问题四：边特征维度不要太高
边特征维度太高，GNN的消息传递会非常慢。我一般控制在10维以内。
3.5 本章小结与知识图谱
好，来总结一下。分子图构建的流程就是：SMILES → RDKit Mol → 提取原子/键特征 → 构建PyG的Data对象。
下面这张图展示了完整的知识体系：
你看，整个过程一环扣一环。SMILES进去，Data对象出来，后面就直接喂给GNN模型了。
我个人觉得，这一步是图神经网络做ADMET预测的基石。地基打不好，后面再高大上的模型也白搭。