字节 Seed AI4S 核心成员离职：AI 赋能科学探索的路线走到哪了？

# 依赖安装：pip install rdkit-pypi torch
# 本示例用 RDKit 处理分子表示，用极简 MLP 预测分子极性（LogP）
# 实际生产中会替换为图神经网络（如 GIN、SchNet）

from rdkit import Chem
from rdkit.Chem import Descriptors
import torch
import torch.nn as nn

# ---- 第一步：从 SMILES 提取分子特征 ----
def mol_to_features(smiles: str) -> torch.Tensor:
    """将 SMILES 转为 RDKit 分子对象，提取基础描述符作为特征向量"""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"无法解析 SMILES: {smiles}")
    features = [
        Descriptors.MolWt(mol),          # 分子量
        Descriptors.NumHDonors(mol),     # 氢键供体数
        Descriptors.NumHAcceptors(mol),  # 氢键受体数
        Descriptors.TPSA(mol),           # 极性表面积
        Descriptors.NumRotatableBonds(mol),  # 可旋转键数
        Descriptors.RingCount(mol),      # 环数
    ]
    return torch.tensor(features, dtype=torch.float32)

# ---- 第二步：极简预测模型 ----
class MolPropertyPredictor(nn.Module):
    """6 维特征 → LogP 预测的 MLP，仅作演示"""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(6, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

# ---- 第三步：跑一个真实分子 ----
model = MolPropertyPredictor()
model.eval()

# 阿司匹林的 SMILES
aspirin_smiles = "CC(=O)OC1=CC=CC=C1C(=O)O"
features = mol_to_features(aspirin_smiles)

with torch.no_grad():
    predicted_logp = model(features)

# 用 RDKit 的真实值做对照
true_logp = Descriptors.MolLogP(Chem.MolFromSmiles(aspirin_smiles))

print(f"阿司匹林 SMILES: {aspirin_smiles}")
print(f"提取特征向量: {features.tolist()}")
print(f"模型预测 LogP: {predicted_logp.item():.2f}")
print(f"RDKit 计算真实 LogP: {true_logp:.2f}")
print(f"（模型未训练，预测值无意义；训练后应逼近真实值）")

# 1. 搭建基础环境（以 conda 为例）
conda create -n ai4s python=3.10
conda activate ai4s
pip install rdkit-pypi torch dgllife  # dgllife 是 DGL 的生命科学工具包

# 2. 获取公开分子数据集（MoleculeNet 是标准基准）
# 下载数据集后用 dgl 的内置 loader 读取
python -c "
from dgllife.data import MoleculeNet
dataset = MoleculeNet('ESOL')  # 水溶性预测数据集
print(f'样本数: {len(dataset)}')
print(f'第一个样本: SMILES={dataset[0][0]}, 标签={dataset[0][1]}')
"

# 3. 跑一个图神经网络基线（约 30 分钟可完成训练）
# 参考 dgllife 官方 example：graph_predictor.py
# 修改数据集名称为 ESOL 或 HIV，即可得到可运行的基线模型