人工智能辅助药物研发的现状与未来：从算法到临床的跨越

# 1. 下载并解压
wget http://bits.csb.pitt.edu/files/crossdock2020.tgz
tar -xzf crossdock2020.tgz

# 2. 过滤分辨率 < 2.5 Å 且口袋-配体 RMSD < 1 Å
python filter_crossdock.py \
  --input_dir crossdock2020 \
  --output_dir crossdock_filtered \
  --max_rmsd 1.0

# models/pocket2mol.py
import torch
from torch_geometric.nn import MessagePassing
from e3nn import o3

class Pocket2MolLayer(MessagePassing):
    def __init__(self, irreps_node, irreps_edge):
        super().__init__(aggr='mean')
        self.tp = o3.FullyConnectedTensorProduct(
            irreps_node, irreps_edge, irreps_node
        )
    def forward(self, x, edge_index, edge_attr):
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)

class Pocket2Mol(torch.nn.Module):
    def __init__(self, num_atom_type=23, hidden_dim=128):
        super().__init__()
        self.embed = torch.nn.Embedding(num_atom_type, hidden_dim)
        self.layers = torch.nn.ModuleList([
            Pocket2MolLayer(hidden_dim, hidden_dim) for _ in range(6)
        ])
        self.atom_head = torch.nn.Linear(hidden_dim, num_atom_type)
        self.pos_head   = torch.nn.Linear(hidden_dim, 3)

    def forward(self, z, pos, edge_index):
        h = self.embed(z)
        for layer in self.layers:
            h = layer(h, edge_index, pos)
        atom_logits = self.atom_head(h)
        delta_pos   = torch.tanh(self.pos_head(h))
        return atom_logits, delta_pos

# train.py
from torch_geometric.loader import DataLoader
from models.pocket2mol import Pocket2Mol
from dataset import CrossDockDataset

train_set = CrossDockDataset("crossdock_filtered/train")
loader = DataLoader(train_set, batch_size=8, shuffle=True)

model = Pocket2Mol().cuda()
optimizer = torch.optim.AdamW(model.parameters(), 1e-3)

for epoch in range(200):
    for data in loader:
        data = data.cuda()
        atom_logits, delta_pos = model(data.z, data.pos, data.edge_index)
        loss_atom = torch.nn.functional.cross_entropy(
            atom_logits, data.y_atom)
        loss_pos  = torch.nn.functional.mse_loss(
            delta_pos, data.y_pos)
        loss = loss_atom + 5 * loss_pos
        optimizer.zero_grad(); loss.backward(); optimizer.step()

# generate.py
from rdkit import Chem
from utils import xyz2mol

model.eval()
with torch.no_grad():
    atom_logits, delta_pos = model(z, pos, edge_index)
    atom_types = torch.argmax(atom_logits, dim=-1)
    new_pos = pos + delta_pos
    mol = xyz2mol(atom_types.cpu().numpy(),
                  new_pos.cpu().numpy())
    Chem.MolToMolFile(mol, 'generated_ligand.mol')

# 使用 GNINA 进行对接
gnina -r receptor.pdb -l generated_ligand.mol \
      --autobox_ligand reference_ligand.mol \
      --score_only  # 快速评估