当前位置：首页 > news >正文

别再只调包了！手把手教你用RDKit和PyTorch Geometric从SMILES字符串构建分子图数据

news 2026/6/11 2:27:56

从SMILES到分子图：用RDKit和PyTorch Geometric构建GNN输入数据的完整指南

当我们在MoleculeNet中加载ESOL数据集时，那些神秘的x、edge_index和edge_attr张量从何而来？本文将带你深入分子图数据的构建过程，理解如何从简单的SMILES字符串生成图神经网络所需的复杂数据结构。

1. 分子图数据的基本概念

分子图是化学信息学中表示分子结构的一种强大方式。在这种表示中，原子作为图中的节点，化学键作为边。PyTorch Geometric（PyG）作为图神经网络的主流框架，需要三种核心数据：

节点特征矩阵(x): 描述每个原子的属性
边索引(edge_index): 定义原子间的连接关系
边属性(edge_attr): 描述化学键的特征

以ESOL数据集中的2-pyrrolidone分子(O=C1CCCN1)为例，其SMILES字符串简洁地描述了分子结构，但要将它转换为图数据，我们需要解决几个关键问题：

如何解析SMILES字符串获取原子和键的信息？
如何选择有意义的原子和键特征？
如何将这些信息编码为PyG可处理的张量格式？

2. 使用RDKit解析SMILES字符串

RDKit是化学信息学中最常用的工具包之一，它能将SMILES字符串转换为丰富的分子对象。让我们从基础解析开始：

from rdkit import Chem smiles = 'O=C1CCCN1' # 2-pyrrolidone的SMILES表示 mol = Chem.MolFromSmiles(smiles)

成功创建分子对象后，我们可以提取详细的原子和键信息。以下是原子级别的特征提取示例：

atom = mol.GetAtomWithIdx(0) # 获取第一个原子(O) print(f"原子序数: {atom.GetAtomicNum()}") print(f"杂化类型: {atom.GetHybridization()}") print(f"是否在环中: {atom.IsInRing()}")

常见的原子特征包括：

特征名称	RDKit获取方法	物理意义
原子序数	GetAtomicNum()	元素类型
杂化状态	GetHybridization()	原子轨道杂化方式
形式电荷	GetFormalCharge()	原子携带的电荷
芳香性	GetIsAromatic()	是否参与芳香体系
氢原子数	GetTotalNumHs()	连接的氢原子数量

3. 构建节点特征矩阵(x)

节点特征矩阵是GNN理解分子结构的基础。我们需要将RDKit提取的原子特征编码为数值形式。以下是一个完整的特征提取和编码流程：

import torch from rdkit import Chem from rdkit.Chem.rdchem import HybridizationType def atom_features(atom): # 原子类型one-hot编码 atomic_num = [atom.GetAtomicNum()] # 杂化类型编码 hybridization = [ int(atom.GetHybridization() == x) for x in (HybridizationType.SP, HybridizationType.SP2, HybridizationType.SP3) ] # 其他原子特征 features = atomic_num + [ atom.GetTotalDegree(), atom.GetFormalCharge(), atom.GetTotalNumHs(), int(atom.GetIsAromatic()), int(atom.IsInRing()) ] + hybridization return torch.tensor(features, dtype=torch.float) # 为分子中所有原子构建特征矩阵 x = torch.stack([atom_features(atom) for atom in mol.GetAtoms()])

这样构建的x矩阵每一行对应一个原子，列对应不同的特征。在实际应用中，你可能需要调整特征选择和编码方式以适应特定任务。

4. 构建边索引和边属性(edge_index & edge_attr)

分子中的化学键构成了图的边结构。我们需要提取两种信息：

连接关系(edge_index): 哪些原子是相连的
键属性(edge_attr): 这些键的类型和特征

def bond_features(bond): bt = bond.GetBondType() return torch.tensor([ bt == Chem.rdchem.BondType.SINGLE, bt == Chem.rdchem.BondType.DOUBLE, bt == Chem.rdchem.BondType.TRIPLE, bt == Chem.rdchem.BondType.AROMATIC, bond.GetIsConjugated(), bond.IsInRing() ], dtype=torch.float) # 构建边索引和边属性 edge_indices = [] edge_attrs = [] for bond in mol.GetBonds(): i = bond.GetBeginAtomIdx() j = bond.GetEndAtomIdx() # 添加双向边(无向图) edge_indices += [[i, j], [j, i]] edge_attrs += [bond_features(bond), bond_features(bond)] edge_index = torch.tensor(edge_indices).t().contiguous() edge_attr = torch.stack(edge_attrs) if edge_attrs else torch.zeros((0, 6))

键特征通常包括：

键类型(单键、双键、三键、芳香键)
是否共轭
是否在环中
立体化学信息(如有需要)

5. 整合为PyG数据对象

将所有组件整合到PyG的Data对象中：

from torch_geometric.data import Data graph_data = Data( x=x, edge_index=edge_index, edge_attr=edge_attr, smiles=smiles, y=torch.tensor([1.07]) # 示例水溶性值 )

现在，这个数据对象可以直接用于PyG的图神经网络模型。完整的转换流程可以封装为一个可复用的函数：

def smiles_to_graph(smiles, y=None): mol = Chem.MolFromSmiles(smiles) if mol is None: return None # 原子特征 x = torch.stack([atom_features(atom) for atom in mol.GetAtoms()]) # 键特征 edge_indices = [] edge_attrs = [] for bond in mol.GetBonds(): i = bond.GetBeginAtomIdx() j = bond.GetEndAtomIdx() edge_indices += [[i, j], [j, i]] edge_attrs += [bond_features(bond)] * 2 edge_index = torch.tensor(edge_indices).t().contiguous() edge_attr = torch.stack(edge_attrs) if edge_attrs else torch.zeros((0, 6)) return Data( x=x, edge_index=edge_index, edge_attr=edge_attr, smiles=smiles, y=torch.tensor([y]) if y is not None else None )

6. 高级技巧与优化建议

在实际应用中，你可能需要考虑以下进阶问题：

特征工程优化

添加更多原子特征：手性、同位素、自由基状态等
考虑空间坐标(如有3D结构)
引入分子指纹作为全局特征

性能优化

批量处理分子数据
使用缓存避免重复计算
并行化特征提取过程

常见问题处理

无效SMILES字符串的容错处理
氢原子的显式/隐式表示选择
不同数据集的特征一致性

# 批量处理示例 from torch_geometric.data import Batch smiles_list = ['O=C1CCCN1', 'CCO', 'c1ccccc1'] # 示例分子 graphs = [smiles_to_graph(s) for s in smiles_list] batch = Batch.from_data_list([g for g in graphs if g is not None])