基于电子等排体的3D分子生成模型 ShEPhERD - 评测

脚本最后还简单地可视化生成的分子。总的来说，就是处理了生成的分子，进行验证，并显示有效的分子结构。下面是具体的工作流程。

遍历生成的 5 个分子，sample_dict 包含每个分子的详细数据。

for b, sample_dict in enumerate(generated_samples):

构建 XYZ 字符串。从每个分子的字典中提取原子信息（原子符号和位置），并构建一个符合 XYZ 格式的字符串。XYZ 格式是一种常见的原子坐标表示方式。

xyz = '' 
x_ = sample_dict['x1']['atoms']
pos_ = sample_dict['x1']['positions']
xyz += f'{len(x_)}\n{b+1}\n'
for a in range(len(x_)):
    atomic_number_ = int(x_[a])
    position_ = pos_[a]
    xyz += f'{rdkit.Chem.Atom(atomic_number_).GetSymbol()} {str(position_[0].round(3))} {str(position_[1].round(3))} {str(position_[2].round(3))}\n'
xyz += '\n'

从 XYZ 字符串创建分子。尝试通过 RDKit 的 MolFromXYZBlock 方法，从构建的 XYZ 字符串中创建一个分子对象，含电荷。如果创建失败（例如格式错误或无效的原子坐标），它会捕捉到异常并跳过当前样本。

try:
    mol_ = rdkit.Chem.MolFromXYZBlock(xyz)
except Exception as e:
    mol_ = None
    print(f'invalid molecule: {e}')
    continue

for c in [0, 1, -1, 2, -2]:
    mol__ = deepcopy(mol_)
    try:
        rdkit.Chem.rdDetermineBonds.DetermineBonds(mol__, charge=c, embedChiral=True)
    except:
        continue
    if mol__ is not None:
        print(c)
        break 

验证分子，跳过不正常的分子。

try:
    assert sum([a.GetNumRadicalElectrons() for a in mol_.GetAtoms()]) == 0, 'has radical electrons'
    mol_.UpdatePropertyCache()
    rdkit.Chem.GetSymmSSSR(mol_)
except Exception as e:
    mol_ = None
    print(f'invalid molecule: {e}')
    continue

显示分子。代码将有效的分子转换为 SMILES（简化分子输入线性表达式）字符串，并显示该分子。

display(rdkit.Chem.MolFromSmiles(rdkit.Chem.MolToSmiles(mol_)))

在生成的5个分子中包含合法有效的 2 个分子结构如下：

有效分子对应的3D结构如下：

3.4.2 条件生成场景

根据项目提供的 RUNME_conditional_generation_MOSESaq.ipynb 文件查看条件场景的生成流程。在条件生成场景下，首先对目标相互作用特征（形状、电势和药效团）进行前向噪声化，得到 t 时刻的状态。在反向去噪过程中，将 ShEPhERD 去噪的状态替换为噪声化的目标。

导入依赖库。

import open3d 
from shepherd_score_utils.generate_point_cloud import (
    get_atom_coords, 
    get_atomic_vdw_radii, 
    get_molecular_surface,
    get_electrostatics,
    get_electrostatics_given_point_charges,
)
from shepherd_score_utils.pharm_utils.pharmacophore import get_pharmacophores
from shepherd_score_utils.conformer_generation import update_mol_coordinates

print('importing rdkit')
import rdkit
from rdkit.Chem import rdDetermineBonds

import numpy as np
import matplotlib.pyplot as plt

print('importing torch')
import torch
import torch_geometric
from torch_geometric.nn import radius_graph
import torch_scatter

import pickle
from copy import deepcopy
import os
import multiprocessing
from tqdm import tqdm

import sys
sys.path.insert(-1, "model/")
sys.path.insert(-1, "model/equiformer_v2")

print('importing lightning')
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import CSVLogger

from lightning_module import LightningModule
from datasets import HeteroDataset

import importlib

from inference import *

选择文章中用来评估的 checkpoint 和基础配置：

chkpt = 'shepherd_chkpts/x1x3x4_diffusion_mosesaq_20240824_submission.ckpt' # checkpoint used for evaluations in preprint

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

model_pl = LightningModule.load_from_checkpoint(chkpt)
params = model_pl.params
model_pl.to(device)
model_pl.model.device = device

该场景下的生成测试基于项目提供的 conformers/np/molblock_charges_NPs.pkl 文件，该文件里包括 3 个分子三维结构和电荷信息。有以下方式可以提取，静电势表面、表面、以及药效团：

1. 基于天然产物的相互作用特征进行条件化

接下来，从项目提供的分子信息文件 conformers/np/molblock_charges_NPs.pkl 中加载分子数据，然后进行处理和计算，以便提取分子特征，比如电荷、分子表面、药效团（pharmacophores）和电场等信息。

加载数据 (pickle 读取)。打开一个名为 molblock_charges_NPs.pkl 的文件，并将其中的数据加载到 molblocks_and_charges 变量中。该文件保存了分子结构和电荷数据，并且可以存储多个分子（该文件包含 3 个分子）。molblocks_and_charges 是一个列表，其中每个元素包含分子的 MolBlock 数据和电荷数据。

mol = rdkit.Chem.MolFromMolBlock(molblocks_and_charges[index][0], removeHs=False)  # target natural product
charges = np.array(molblocks_and_charges[index][1])  # xTB partial charges in implicit water

with open('conformers/np/molblock_charges_NPs.pkl', 'rb') as f:
    molblocks_and_charges = pickle.load(f)  # len(molblocks_and_charges) == 3

选择目标天然产物。这里，index 选择了第一个天然产物（分子）。index 也可以是 1 或 2，表示选择不同的分子。

index = 0  # 0, 1, 2

从 MolBlock 生成分子对象 (rdkit)。使用 rdkit.Chem.MolFromMolBlock 从 MolBlock 格式的数据中构建一个分子对象 (mol)。charges 是从 molblocks_and_charges 中提取的电荷数据，转换为 NumPy 数组，表示的是在隐式水环境下的部分电荷。

mol = rdkit.Chem.MolFromMolBlock(molblocks_and_charges[index][0], removeHs=False)  # target natural product
charges = np.array(molblocks_and_charges[index][1])  # xTB partial charges in implicit water

选择的天然产物分子的结构如下：

提取分子坐标并进行中心化。mol.GetConformer().GetPositions() 获取分子的三维坐标。mol_coordinates - np.mean(mol_coordinates, axis=0) 将分子坐标进行中心化，即将分子坐标的质心移到原点（这样可以减少坐标的偏移对后续计算的影响）。update_mol_coordinates 用来更新分子对象的坐标。

mol_coordinates = np.array(mol.GetConformer().GetPositions())
mol_coordinates = mol_coordinates - np.mean(mol_coordinates, axis=0)
mol = update_mol_coordinates(mol, mol_coordinates)

获取分子相关的条件性目标。centers 是分子中各个原子的三维坐标。get_atomic_vdw_radii(mol) 获取分子中各个原子的范德华半径（Van der Waals radii）。get_molecular_surface 计算分子的表面，使用了范德华半径、表面采样点的数量、探针半径等参数。

centers = mol.GetConformer().GetPositions()
radii = get_atomic_vdw_radii(mol)
surface = get_molecular_surface(
    centers, 
    radii, 
    params['dataset']['x3']['num_points'], 
    probe_radius=params['dataset']['probe_radius'],
    num_samples_per_atom=20,
)

获取药效团特征。get_pharmacophores 函数用于从分子中提取药效团信息。multi_vector 和 check_access 是通过参数 params 提供的，用来控制药效团计算的不同细节。

pharm_types, pharm_pos, pharm_direction = get_pharmacophores(
    mol,
    multi_vector=params['dataset']['x4']['multivectors'],
    check_access=params['dataset']['x4']['check_accessibility'],
)

计算电场（静电势）。get_electrostatics_given_point_charges 计算给定电荷分布下的静电场（电势），这里使用了 charges（电荷）、centers（原子坐标）和 surface（分子表面）来计算。

electrostatics = get_electrostatics_given_point_charges(
    charges, centers, surface,
)

2. 基于 PDB 配体的相互作用特征进行条件化

这段代码与前一段代码非常相似，主要区别在于它加载的是 PDB（蛋白质数据库）格式的分子数据，而不是之前的天然产物的 MolBlock 数据。下面会逐步解释这段代码。

加载 PDB 数据 (pickle 读取)。打开一个名为 molblock_charges_pdb_lowestenergy.pkl 的文件，并将其中的数据加载到 molblocks_and_charges 变量中。这个文件包含了 PDB 配体（ligand）及其电荷信息。

with open('conformers/pdb/molblock_charges_pdb_lowestenergy.pkl', 'rb') as f:
    molblocks_and_charges = pickle.load(f)

选择目标 PDB 配体。在这里，index 被设定为 6，表示从 molblocks_and_charges 中选择第七个配体。index 可以是 0 到 6 之间的任何值，表示选择不同的配体。随后，提取电荷、并计算表面、静电势表面：

index = 6 

mol = rdkit.Chem.MolFromMolBlock(molblocks_and_charges[index][0], removeHs=False)  # target natural product
charges = np.array(molblocks_and_charges[index][1])  # xTB partial charges in implicit water

分子结构如下：

提取分子坐标并进行中心化。mol.GetConformer().GetPositions() 获取分子的三维坐标。通过 mol_coordinates - np.mean(mol_coordinates, axis=0) 将分子坐标进行中心化，使得分子的质心处于原点。update_mol_coordinates(mol, mol_coordinates) 更新分子对象的坐标。

mol_coordinates = np.array(mol.GetConformer().GetPositions())
mol_coordinates = mol_coordinates - np.mean(mol_coordinates, axis=0)
mol = update_mol_coordinates(mol, mol_coordinates)

获取分子相关的条件性目标。centers 变量保存了分子中各个原子的三维坐标。get_atomic_vdw_radii(mol) 获取分子中各个原子的范德华半径（Van der Waals radii）。get_molecular_surface 计算分子的表面，并根据分子的原子坐标、范德华半径以及其他参数（如探针半径、采样点数量等）来确定表面。

centers = mol.GetConformer().GetPositions()
radii = get_atomic_vdw_radii(mol)
surface = get_molecular_surface(
    centers, 
    radii, 
    params['dataset']['x3']['num_points'], 
    probe_radius=params['dataset']['probe_radius'],
    num_samples_per_atom=20,
)

获取药效团特征。get_pharmacophores 提取分子的药效团特征。multi_vector 和 check_access 是从 params 中获取的参数，控制药效团计算的细节，如药效团的类型和分子是否可接近。

pharm_types, pharm_pos, pharm_direction = get_pharmacophores(
    mol,
    multi_vector=params['dataset']['x4']['multivectors'],
    check_access=params['dataset']['x4']['check_accessibility'],
)

electrostatics = get_electrostatics_given_point_charges(
    charges, centers, surface,
)

由于篇幅原因，另一种方式就略过了。

3.4.3 通过 inpainting 进行条件生成

前面介绍了不同的应用场景怎么来获取相互作用特征，接着介绍根据这些特征来进行条件生成的流程。这部分调用了一个函数 inference_sample() 来生成分子样本或进行相关的推理（inference）任务。具体来说，函数的参数配置了不同的生成选项，其中包含了与分子坐标、药效团、表面、电荷等多个特征相关的设置。下面是对每个部分的详细解释。

设置生成分子的原子数目，batch_size 以及药效团数目

n_atoms = 70
batch_size = 2 # 24 GB
num_pharmacophores = len(pharm_types) # must equal pharm_pos.shape[0] if inpainting

调用 inference_sample 函数进行推理生成分子。调用 inference_sample 生成分子样本的代码如下：

generated_samples = inference_sample(
    model_pl,
    batch_size = batch_size,
    
    N_x1 = n_atoms,
    N_x4 = num_pharmacophores,
    
    unconditional = False,
    
    prior_noise_scale = 1.0,
    denoising_noise_scale = 1.0,
    
    inject_noise_at_ts = [],
    inject_noise_scales = [],    
    
    harmonize = False,
    harmonize_ts = [],
    harmonize_jumps = [],
    
    
    # all the below options are only relevant if unconditional is False
    
    inpaint_x2_pos = False, # note that x2 is implicitly modeled via x3
    
    inpaint_x3_pos = True,
    inpaint_x3_x = True,
    
    inpaint_x4_pos = True,
    inpaint_x4_direction = True,
    inpaint_x4_type = True,
    
    stop_inpainting_at_time_x2 = 0.0,
    add_noise_to_inpainted_x2_pos = 0.0,
    
    stop_inpainting_at_time_x3 = 0.0,
    add_noise_to_inpainted_x3_pos = 0.0,
    add_noise_to_inpainted_x3_x = 0.0,
    
    stop_inpainting_at_time_x4 = 0.0,
    add_noise_to_inpainted_x4_pos = 0.0,
    add_noise_to_inpainted_x4_direction = 0.0,
    add_noise_to_inpainted_x4_type = 0.0,
    
    # these are the inpainting targets
    center_of_mass = np.zeros(3), # center of mass of x1; already centered to zero above
    surface = surface,
    electrostatics = electrostatics,
    pharm_types = pharm_types,
    pharm_pos = pharm_pos,
    pharm_direction = pharm_direction,
)

batch_size 设置为 2，显存占用大约为 10 GB，采样过程耗时约 2 分钟。采样生成的结构信息保存在 generated_samples 变量中。generated_samples 的类型是 List[Dict] ，列表中的每个字典包含四个键 x1, x2, x3, x4，分别表示分子结构、表面几何、表面电荷/药效团特征。

类似于无条件生成场景，脚本最后也简单地可视化生成的分子。总的来说，就是处理了生成的分子，进行验证，并显示有效的分子结构。

# quick visualization of generated samples
# full analyses, including extensive validity checks, can be performed by following https://github.com/coleygroup/shepherd-score

for b,sample_dict in enumerate(generated_samples):
    
    xyz = '' 
    
    x_ = sample_dict['x1']['atoms']
    pos_ = sample_dict['x1']['positions']
    
    xyz += f'{len(x_)}\n{b+1}\n'
    for a in range(len(x_)):
        atomic_number_ = int(x_[a])
        position_ = pos_[a]
        
        xyz+= f'{rdkit.Chem.Atom(atomic_number_).GetSymbol()} {str(position_[0].round(3))} {str(position_[1].round(3))} {str(position_[2].round(3))}\n'
    xyz+= '\n'
    
    try:
        mol_ = rdkit.Chem.MolFromXYZBlock(xyz)
    except Exception as e:
        mol_ = None
        print(f'invalid molecule: {e}')
        continue
    
    try:
        for c in [0, 1, -1, 2, -2]:
            mol__ = deepcopy(mol_)
            try:
                rdkit.Chem.rdDetermineBonds.DetermineBonds(mol__, charge = c, embedChiral = True)
            except:
                continue
            if mol__ is not None:
                print(c)
                break 
    except Exception as e:
        mol_ = None
        print(f'invalid molecule: {e}')
        continue
    
    mol_ = mol__
    try:
        assert sum([a.GetNumRadicalElectrons() for a in mol_.GetAtoms()]) == 0, 'has radical electrons'
        mol_.UpdatePropertyCache()
        rdkit.Chem.GetSymmSSSR(mol_)
        
    except Exception as e:
        mol_ = None
        print(f'invalid molecule: {e}')
        continue

    display(rdkit.Chem.MolFromSmiles(rdkit.Chem.MolToSmiles(mol_)))
    
    continue

生成的两个分子都是有效分子，结构如下所示：

对应的两个有效分子的 3D 结构如下：

3.5 自定义案例（3WZE）条件分子生成

作为测试，是我们的老演员，3WZE。不同的是，我们选择 3WZE 同一蛋白的不同小分子体系 2QU5 也作为输入。我们的计划是，使用3WZE小分子的药效团、2QU5小分子的表面和电荷，进行分子条件生成。

3WZ E和 2QU5的align结果见下图，可以看到小分子的叠合非常好，分子结构也非常相似，只有骨架中间部分的差异。

分子生成过程详见 RUNME_conditional_generation_3WZE.ipynb 文件，这里就略过。在3090显卡上，生成120个分子耗时1个小时，实际上有效的分子只有82个。

随后，使用类似 DrugHIVE 类似的qvina对接流程，将生成的小分子对接到 3WZE口袋中。具体过程参考 ./3WZE/get_qvina， 这里不再赘述。

以下是生成分子按照对接打分由高到低的结果。大部分的分子是非常类药的，结构也比较合理：

Top 3的分子结构如下，对接打分分别为：-11.8，-11.1， 11.0。（但是第一个分子结构挺奇怪的，可能是因为化学键重建的时候，原子距离有点异常导致的）

从上述结果来看，ShEPhERD 模型生成的分子还是不错的，值得一试。