• 梯度引导的分子生成扩散模型- GaUDI 评测

    GaUDI模型来自于以色列理工Tomer Weiss的2023年发表在预印本ChemRxiv上的工作 《Guided Diffusion for Inverse Molecular Design》。原文链接：Guided Diffusion for Inverse Molecular Design | Materials Chemistry | ChemRxiv | Cambridge Open Engage

    GaUDI模型有点像强化学习，目的都是生成特定性质的分子，且效果均非常显著，而且生成过程均存在迭代更新过程。

    写在前面

    GaUDI模型有点像强化学习，目的都是生成特定性质的分子，且效果均非常显著，而且生成过程均存在迭代更新过程。

    但是，二者原理不同。

    从迭代过程来说，强化学习生成特定属性的分子迭代发生在生成的批次分子之间，而梯度引导扩散模型是发生在生成一个分子的过程中，因此，从速度上来说，GaUDI模型非常快，很有优势。

    从目标函数的形式来说，强化学习是使用目标函数计算生成的分子，获得reward。reward是 一个数值，而是通过数值，调节分子生成器的梯度强度，让分子生成器降低目标分子的损失，于是分子生成器生成更多特定属性的分子。GaUDI模型则不同，目标函数产生的梯度，并未返回到分子生成器中，而是直接用于引导Zt-1，使Zt-1亲向特定属性（方向）。因此，对于任何的分子性质要求，GaUDI模型的分子生成器，即扩散模型，都可以通用，无需重新训练。

    从联合概率分布来说，GaUDI模型是联合概率分布，而强化学习不是。GaUDI模型生成的分子更加”平滑“，更具有实用性。强化学习更容易钻模型的漏洞，生成一些奇怪的分子。

    一、背景介绍

    GaUDI是用于逆向分子设计的一个引导扩散模型，因此文章的题目是《Guided Diffusion for Inverse Molecular Design》。GaUDI对扩散模型中的分子生成器进行了改造，将分子生成过程使用性质预测的函数进行梯度引导，更新Zt-1，使Zt-1更靠近我们期待的性质（例如：logP等）。我们期待的性质可以使单目标也可以是多目标，一切取决于我们的梯度函数f(x,t)。

    引导分子生成的梯度来自于目标函数f(x,t), 函数f(x,t)用于根据输入的去噪过程中的Zt-1，预测Zt-1(即：分子)对应的性质。利用预测性质与目标性质之间的损失梯度，更新Zt-1，那么在去噪生成过程中，生成的Zt-1就会越来越靠近我们的期望的性质。当然目标函数f(x,t)需要另外训练。梯度对Zt-1的更新程度可使用调节系数（gradient scalar values, s）来控制。整个去噪采样逻辑如下图：

    伪代码如下：

    二、文章案例

    2.1 多目标优化（全局最优）

    作者以 图 A 的 cc-PBHs 为例，目标是增加改分子的稳定性（更低的Erel）。这里作者是进行全局最优的例子，目标函数是生成分子与cc-PBHs 的LUMO, HLG, IP, EA 的MSE加上Erel。作者尝试了不同的调节系数（gradient scalar values, s）结果如图B所示。图C展示了一些具体的例子。

    2.2 Out-of-distribution生成

    验证GaUDI模型是否可以生成更高的HLG值的分子，限定条件是在相同的环数量下。结果如下图：

    图A是数据集中不同环数下的最高HLG分子，图B是生成的分子，图C是不同调节系数下的结果。从结果上来看，生成分子的效果非常显著。

    2.3 多目标优化

    将目标设置为：ℓ(HLG,IP, EA) = 3 · HLG + IP − EA, 结果如下下图：

    整体结果还是非常惊喜的。因为没有任何的强化学习，但是实现了类似的结果，效果非常明显。非常好的弥补了强化学习只能身材个核那个smiles的缺点。

    原文代码开源：porannegroup / GaUDI · GitLab

    三、模型应用测试

    3.1 环境安装

    首先，复制代码：

    git clone https://github.com/tomer196/GaUDI.git

    代码目录：

    1. .
    2. ├── analyze
    3. ├── cond_prediction
    4. ├── data
    5. ├── edm
    6. ├── environment.yml
    7. ├── eval_validity.py
    8. ├── GaUDI.png
    9. ├── generation_guidance.py
    10. ├── __init__.py
    11. ├── LICENSE
    12. ├── models_edm.py
    13. ├── __pycache__
    14. ├── README.md
    15. ├── sampling_edm.py
    16. ├── train_edm.py
    17. └── utils

    安装环境目录，并激活环境：

    1. conda env create -n GaUDI --file environment.yml
    2. conda activate GaUDI

    如果不行，就分开安装rdkit, pytorch等，注意版本：

    1. conda install pytorch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 pytorch-cuda=11.8 -c pytorch -c nvidia
    2. conda install -c conda-forge rdkit
    3. pip install matplotlib networkx tensorboard pandas scipy tqdm imageio

    安装过程，比较顺利，没有特殊问题。但是centos和win系统比较麻烦，估计要摸索一下。

    另外，需要安装几个附属模块：

    1. pip install tensorboard
    2. pip install imgaug

    3.2 数据准备&下载数据

    作者使用到两个数据集：COMPAS 和 PASS，下载链接：

    COMPAS：porannegroup / COMPAS · GitLab

    PASS：PASS molecular dataset

    需要从上述网址中，下载三个文件，分别是：PASs_csv.tar.gz， PASs_xyz.tar.gz， compas-main-COMPAS-1.zip

    在data目录下，创建一个all_data文件夹，将下载的PASs_csv.tar.gz， PASs_xyz.tar.gz， compas-main-COMPAS-1.zip文件上传。然后，分别解压：

    1. unzip compas-main-COMPAS-1.zip
    2. tar -zxvf PASs_csv.tar.gz
    3. tar -zxvf PASs_xyz.tar.gz
    4.  
    5. #删除压缩文件
    6. rm compas-main-COMPAS-1.zip PASs_csv.tar.gz PASs_xyz.tar.gz
    然后去更改./data/aromatic_dataloaders.py中get_path函数，关于数据的路径。
    1. def get_paths(args):
    2.     if not hasattr(args, "dataset"):
    3.         csv_path = args.csv_file
    4.         xyz_path = args.xyz_root
    5.     elif args.dataset == "cata":
    6.         csv_path = "/home/tomerweiss/PBHs-design/data/COMPAS-1x.csv"
    7.         xyz_path = "/home/tomerweiss/PBHs-design/data/peri-cata-89893-xyz"
    8.     elif args.dataset == "peri":
    9.         csv_path = "/home/tomerweiss/PBHs-design/data/peri-xtb-data-55821.csv"
    10.         xyz_path = "/home/tomerweiss/PBHs-design/data/peri-cata-89893-xyz"
    11.     elif args.dataset == "hetro":
    12.         csv_path = "/home/tomerweiss/PBHs-design/data/db-474K-OPV-filtered.csv"
    13.         xyz_path = "/home/tomerweiss/PBHs-design/data/db-474K-xyz"
    14.     elif args.dataset == "hetro-dft":
    15.         csv_path = "/home/tomerweiss/PBHs-design/data/db-15067-dft.csv"
    16.         xyz_path = ""
    17.     else:
    18.         raise NotImplementedError
    19.     return csv_path, xyz_path

    改为：

    1. def get_paths(args):
    2.     if not hasattr(args, "dataset"):
    3.         csv_path = args.csv_file
    4.         xyz_path = args.xyz_root
    5.     elif args.dataset == "cata":
    6.         csv_path = "./data/all_data/compas-main-COMPAS-1/COMPAS-1/COMPAS-1x.csv"
    7.         xyz_path = "/data/all_data/compas-main-COMPAS-1/COMPAS-1/pahs-cata-34072-xyz"
    8.     elif args.dataset == "hetro":
    9.         csv_path = "./data/all_data/db-474K-OPV-filtered.csv"
    10.         xyz_path = "./data/all_data/db-474K-xyz"
    11.     else:
    12.         raise NotImplementedError
    13.     return csv_path, xyz_path

    3.3 训练EDM 模型

    由于作者没有提供预训练好的checkpoint, 所以需要我们自己训练。回到主目录，然后：

    python train_edm.py

    训练过程了会被记录，将保存在./summary路径下，默认是：cata-test文件夹。整个训练过程，3090显卡需要5h左右。

    训练完成后，./summary/cata-test内得到如下目录：

    1. .
    2. ├── args.txt
    3. ├── epoch_0
    4. ├── epoch_100
    5. ├── epoch_150
    6. ├── epoch_200
    7. ├── epoch_250
    8. ├── epoch_300
    9. ├── epoch_350
    10. ├── epoch_400
    11. ├── epoch_450
    12. ├── epoch_50
    13. ├── epoch_500
    14. ├── epoch_550
    15. ├── epoch_600
    16. ├── epoch_650
    17. ├── epoch_700
    18. ├── epoch_750
    19. ├── epoch_800
    20. ├── epoch_850
    21. ├── epoch_900
    22. ├── epoch_950
    23. ├── events.out.tfevents.1700200741.a01.3529424.0
    24. ├── events.out.tfevents.1700200774.a01.3532044.0
    25. └── model.pt #模型文件

    3.4 训练条件模型

    同样的，作者没有给checkpoint，条件模型也需要训练。在主目录下，执行：

    python cond_prediction/train_cond_predictor.py

    训练过程将保存在prediction_summary目录中的cata-test文件夹中。训练过程大约也需要4h。训练完成后， prediction_summary的目录如下：

    1. .
    2. └── cata-test
    3.     ├── args.txt
    4.     ├── events.out.tfevents.1700211860.a01.63329.0
    5.     └── model.pt

    3.5 条件引导扩散的分子生成

    3.5.1 使用generation_guidance.py

    在完成EDM和条件模型的训练以后，现在可以进行条件引导分子扩散。

    首先，需要修改generation_guidance.py文件中的EDM和条件模型路径。将如下内容：

    1. args = get_edm_args("/home/tomerweiss/PBHs-design/summary/hetro_l9_c196_orientation2")
    2. cond_predictor_args = get_cond_predictor_args(
    3.         f"/home/tomerweiss/PBHs-design/prediction_summary/"
    4.         f"cond_predictor/hetro_gap_homo_ea_ip_stability_polynomial_2_with_norm"
    5.     )

    根据我们前面训练获得的模型，更改为：

    1. args = get_edm_args("./gaudi/summary/cata-test")
    2. cond_predictor_args = get_cond_predictor_args(
    3.         f"./gaudi/prediction_summary/cata-test"
    4.     )

    在189~193行可以修改，需要生成的分子数量和梯度引导强度。

    1. ##################### 设置引导梯度的强度，和想要的分子数量  ###################################
    2. # Set controllable parameters
    3. args.batch_size = 512   # number of molecules to generate
    4. scale = 0.6             # gradient scale - the guidance strength
    5. n_nodes = 10            # number of rings in the generated molecules
    6. #########################################################################################

    然后执行：

    python generation_guidance.py

    就可以生成分子。

    但是，存在的问题是：

    1. plot_graph_of_rings
    2. plot_rdkit

    这两个函数一直无法通过。同时发现，改模型输出的结果不是xyz文件，也不是sdf文件，似乎是图片。这是因为，作者在定义图的时候与我们平常的方法不同，平常的，我们会将一个原子作为一个节点，但是这里作者是将一个环作为一个节点。因此输出的结果比较奇怪。需要其他的转化函数（plot_rdkit），才能输出我们能查看的RDKIT对象。但是我们在运行generation_guidance.py文件时，plot_rdkit文件报错了。所以无法继续运行。

    3.5.2 generation_guidance.py文件修正

    在研究了其代码和输出结果以后，我们发现代码里面存在变量错误和文件名命名错误。

    将以下代码进行修改：

    1. plot_graph_of_rings(
    2.     x_stable[idx].detach().cpu(),
    3.     atom_type_stable[idx].detach().cpu(),
    4.     filename=f"{dir_name}/{i}.pdf",
    5.     title=f"{best_str}\n {value}",
    6.     dataset=args.dataset,
    7.         )
    8.  
    9. plot_rdkit(
    10.     x_stable[idx].detach().cpu(),
    11.     atom_type_stable[idx].detach().cpu(),
    12.     filename=f"{dir_name}/mol_{i}.pdf",
    13.     title=f"{best_str}\n {value}",
    14.     dataset=args.dataset,
    15.         )

    修改为：

    1. try:
    2.   plot_graph_of_rings(
    3.       x_stable[idx].detach().cpu(),
    4.       atom_type[idx].detach().cpu(),
    5.       filename=f"{dir_name}/{i}.png",
    6.       title=f"{best_str}\n {value}",
    7.       dataset=args.dataset,
    8.   )
    9.   plot_rdkit(
    10.       x_stable[idx].detach().cpu(),
    11.       atom_type[idx].detach().cpu(),
    12.       filename=f"{dir_name}/mol_{i}.png",
    13.       title=f"{best_str}\n {value}",
    14.       dataset=args.dataset,
    15.   )
    16.   
    17. except:
    18.    pass

    使用try语句的原因是，及时通过了有效检测，但是仍然有分子是画不出来的。

    生成分子：

    python generation_guidance.py

    运行结束后，将产生./best文件夹目录，有一个以运行日期为名的文件夹，例如：1120_00:42:46_0.6，其中保存对于 Top 5 最优分子的结果，包括节点图以及分子结构。

    生成目录示例：

    1. ./best
    2. └── 1120_00:42:46_0.6
    3.     ├── 0.png
    4.     ├── 1.png
    5.     ├── 2.png
    6.     ├── 3.png
    7.     ├── 4.png
    8.     ├── all.png
    9.     └── mol_0.png

    生成节点图：

    目标函数分子的gap最大，最优分子结构：

    由于生成分子结构代码的问题，最终只生成了一个分子的结构图，如下图所示：

    1. scale=0.8
    2. stability_dict['mol_valid']=82.00% out of 50
    3. Mean target function value: -4.7209
    4. best value: -5.982420921325684, pred: tensor([ 0.2653, -6.9535, -0.1940,  0.3079, 53.0093])
    5. Mean target function value (from valid): -4.7066

    根据输出的结果，最优的分子的目标函数值，gap值为：5.98，但是由于画图的原因，导致无法输出分子结构。能输出分子结构的目标函数值是5.488。

    因为在没有引导的时候，即scale参数为0时（在generation_guidance.py文件内调整scale参数的值），相当于直接使用EDM进行分子生成，运行代码输出：

    1. scale=0
    2. stability_dict['mol_valid']=100.00% out of 50
    3. Mean target function value: 0.0738
    4. best value: -1.7954111099243164, pred: tensor([ 1.1055, -8.5804, -0.2251,  0.3761, 55.9974])
    5. Mean target function value (from valid): 0.0738

    对应的gap的最大值只有1.79，远小于scale为0.8的结果，说明作者提出的梯度引导很符合target函数的要求，效果还不错。但是牺牲的代价是分子的有效率降低了。

    3.6 scale梯度调节参数

    刚才我们提到scale参数，scale参数是调节引导强度的超参数，当这个值越大时，则生成的分子越满足，越靠近我们的目标函数需求，但是分子的有效率会大大降低，例如，我们将scale设置为2时：只有6%的分子有效。

    1. scale=2
    2. stability_dict['mol_valid']=6.00% out of 50
    3. Mean target function value: -5.2808
    4. best value: -6.993871212005615, pred: tensor([ 0.2147, -6.5605, -0.1836,  0.2928, 52.0480])
    5. Mean target function value (from valid): -5.1165
    6. best value (from stable): tensor([ 0.2562, -7.1488, -0.2008,  0.3158, 52.8709]), score: -5.479823112487793
    7. best value (from stable): tensor([ 0.3591, -7.2578, -0.1995,  0.3182, 53.5928]), score: -5.1993408203125
    8. best value (from stable): tensor([ 0.1914, -7.4633, -0.2178,  0.3219, 54.1037]), score: -4.670310020446777

    我们将scale设置为4时，报错了，生成了一些奇怪的分子无法计算其gap等分子属性结果，分子的有效率直接降为0：

    1. scale=4
    2. stability_dict['mol_valid']=0.00% out of 50
    3. Mean target function value: -2.2682
    4. best value: -6.642547130584717, pred: tensor([ 0.3690, -6.6970, -0.1769,  0.3109, 51.9043])
    5. Mean target function value (from valid): nan
    6. ~/anaconda3/envs/GaUDI/lib/python3.8/site-packages/numpy/lib/histograms.py:883: RuntimeWarning: invalid value encountered in divide
    7.   return n/db/n.sum(), bin_edges

    3.7 条件引导扩散的分子生成（jupyter notebook版本）

    作者也提供了一个jupyter notebook的版本，详见文件中的Case_Test.ipynb文件。

    导入相关模块：

    1. import shutil
    2. import os
    3. import matplotlib.pyplot as plt
    4. import torch
    5. import numpy as np
    6. import gdown
    7. from time import time
    8. from types import SimpleNamespace
    9.  
    10. from cond_prediction.train_cond_predictor import get_cond_predictor_model
    11. from data.aromatic_dataloader import create_data_loaders
    12. from models_edm import get_model
    13. from utils.helpers import switch_grad_off, get_edm_args, get_cond_predictor_args

    加载预训练模型：

    edm_model - 扩散模型，一种在每次迭代时稍微清洁分子的降噪器。

    cond_predictor - 条件预测器，预测生成过程中噪声分子的属性。 以时间为条件。

    1. from cond_prediction.train_cond_predictor import get_cond_predictor_model
    2. from data.aromatic_dataloader import create_data_loaders
    3. from models_edm import get_model
    4. from utils.helpers import switch_grad_off, get_edm_args, get_cond_predictor_args
    5.  
    6.  
    7. # 加载EDM 参数配置
    8. args = get_edm_args("pre_trained/EDM")
    9. # 加载条件控制模型参数配置
    10. cond_predictor_args = get_cond_predictor_args("pre_trained/cond_pred")
    11.  
    12.  
    13. # 数据集的模拟，我们仅需要它来了解数据的形状及其统计数据。
    14. dataset_mock = SimpleNamespace( 
    15.     num_node_features=12,
    16.     num_targets=5,
    17.     mean=torch.tensor([ 1.1734, -9.2780, -0.2356,  0.4042, 53.9615]),
    18.     std=torch.tensor([0.5196, 0.3886, 0.0179, 0.0152, 1.6409]),
    19. )
    20. train_loader_mock = SimpleNamespace(dataset=dataset_mock)
    21.  
    22.  
    23. # 加载EDN模型
    24. edm_model, nodes_dist, prop_dist = get_model(args, train_loader_mock, only_norm=True)
    25. # 加载条件控制模型
    26. cond_predictor = get_cond_predictor_model(cond_predictor_args, dataset_mock)
    27. # 设置模型中的参数为不可训练
    28. switch_grad_off([edm_model, cond_predictor])

    定义生成的内容：

    batch_size - 有多少分子 - 影响运行时间。

    n_rings 每个分子中环的数量。

    Guiding_scale - 对具有较低目标功能的分子的引导强度 - 梯度标量影响生成分子的有效性与其目标分数之间的权衡。

    它越高，有效分子的百分比就会减少。 然而，生成的有效分子将更接近定义的目标属性。

    目标函数 - 我们想要最小化的函数。 预测属性的任意组合。

    1. ### 定义分子限制和目标函数 ###
    2. batch_size = 10
    3. n_rings = 8
    4. guidance_scale = 1
    5.  
    6. def target_function(_input, _node_mask, _edge_mask, _t):
    7.     pred = cond_predictor(_input, _node_mask, _edge_mask, _t)
    8.     pred = prop_dist.unnormalize(pred)
    9.     gap = pred[:, 0]
    10.     homo = pred[:, 1]
    11.     ea = pred[:, 2] * 27.2114  # hartree to eV
    12.     ip = pred[:, 3] * 27.2114  # hartree to eV
    13.     return -gap
    14.     # return ip + ea + 3 * gap

    生成分子：

    1. from sampling_edm import sample_guidance
    2.  
    3. # 初始化生成分子，由于这里的模型是一个节点是一个环，所以节点比较奇怪
    4. nodesxsample = torch.Tensor([n_rings] * batch_size).long()
    5. # nodesxsample = nodes_dist.sample(args.batch_size)
    6.  
    7. start_time = time()
    8. x, one_hot, node_mask, edge_mask = sample_guidance(
    9.     args,
    10.     edm_model,
    11.     target_function,
    12.     nodesxsample,
    13.     scale=guidance_scale,
    14. )
    15. print(f"Generated {x.shape[0]} molecules in {time() - start_time:.2f} seconds")

    使用 RDKit 评估分子的有效性并过滤掉无效分子：

    1. target_function_values = get_target_function_values(
    2.       x, one_hot, target_function, node_mask, edge_mask, edm_model
    3.   )
    4. pred = prop_dist.unnormalize(predict(cond_predictor, x, one_hot, node_mask, edge_mask, edm_model))

    绘制最佳分子（具有较低的目标函数值）。 生成的分子表示为环图 (GOR)，并带有表示环方向的加法节点。

    1. from utils.plotting import plot_grap_of_rings_inner, plot_rdkit
    2.  
    3. best_idxs = target_function_values.argsort()
    4. n_plot = min(3, x.shape[0])
    5. fig, axes = plt.subplots(1, n_plot, figsize=(5*n_plot, 7))
    6. for i, idx in enumerate(best_idxs[:n_plot]):
    7.     pred_i = pred[idx]
    8.     title = f"HLG: {pred_i[0]:.2f}, HOMO: {pred_i[1]:.2f}, EA: {pred_i[2]:.2f}, IP: {pred_i[3]:.2f}"
    9.     plot_grap_of_rings_inner(axes[i], x[idx], one_hot.argmax(2)[idx], title=title, dataset=args.dataset, )
    10. plt.show()

    输出：

    转换成完整的分子结构：

    1. fig, axes = plt.subplots(1, n_plot, figsize=(5*n_plot, 7))
    2. for i, idx in enumerate(best_idxs[:n_plot]):
    3.     pred_i = pred[idx]
    4.     title = f"HLG: {pred_i[0]:.2f}, HOMO: {pred_i[1]:.2f}, EA: {pred_i[2]:.2f}, IP: {pred_i[3]:.2f}"
    5.     plot_rdkit(x[idx], one_hot.argmax(2)[idx], axes[i], title=title, dataset=args.dataset)
    6. plt.show()

    四、总结

    1. GaUDI是一个梯度引导的分子生成扩散模型，可以在不重新训练分子生成模型的基础上，通过目标函数的梯度引导分子生成，从而产生特定属性的分子。文中以Gap等分子属性作为目标，验证了其方法的可行性。

    2. 对于药物设计，分子设计等而言，该方法不能直接使用。首先，该模型对分子的表示，以环为节点，不符合我们几何神经网络中关于分子的设定，导致分子还原存在困难。这也是为什么上述结果中，有着GOA图的原因，以及有大量分子无法输出结构的原因。数据集中的分子均为多环结构，不是药物设计中常用的分子结构，模型的偏差很严重。

    结论是，该模型不推荐，难用不友好，但是该方法提供了非常好的概念验证，值得进一步研究学习。

  • 相关阅读:
    css 三栏布局的实现？
    数据结构 - 数组
    期货开户手续费是怎么查询？
    【探索Linux】—— 强大的命令行工具 P.26（网络编程套接字基本概念—— socket编程接口 | socket编程接口相关函数详细介绍 ）
    2016-2023全国MEM国家A类线趋势图：浙大MEM要高多少？
    unity打AB包，AssetBundle预制体与图集（一）
    Sandboxie+Buster Sandbox Analyzer打造个人沙箱
    java中获取主机信息InetAddress类的两种方法
    LocalDate、LocalTime、LocalDateTime常用方法
    controller调用service层报错Invalid bound statement (not found)
  • 原文地址：https://blog.csdn.net/wufeil7/article/details/134541056