相互作用的检索增强 3D 分子生成扩散模型

IRDiff 是一个全新的基于蛋白质-配体相互作用的检索增强 3D 分子扩散模型，可以生成目标感知的分子。IRDiff 利用一组设计好的参考配体分子来引导扩散模型生成满足目标特性的分子。

一、背景介绍

IRDiff 来源于清华大学深圳国际研究生院的杨文明教授和鹏城实验室的王宇研究员为通讯作者的文章：《Interaction-based Retrieval-augmented Diffusion Models for Protein-specific 3D Molecule Generation》。IRDiff 是 Interaction-based Retrieval-augmented Diffusion model 的缩写。文章链接：Interaction-based Retrieval-augmented Diffusion Models for Protein-specific 3D Molecule Generation | OpenReview。该文章于 2024 年 6 月 25 日发表在 ICML 2024 会议上。

很多基于结构的分子生成模型采用从头（from scratch）生成的策略，不使用参考分子可能导致模型不能实现最理想的结果。为了解决这个问题，作者提出了 IRDiff 这一基于相互作用的检索增强 3D 分子扩散模型，以实现目标感知的分子生成。

IRDiff 利用了一个精心挑选的配体参考集，其中的配体都是具有所需特性的（如高结合亲和力），通过这一参考集引导扩散模型，合成满足设计标准的配体。具体来说，作者设计了一个几何蛋白质-分子相互作用网络（PMINet），并使用结合亲和力信号对其进行预训练，以： (i) 检索对目标蛋白质具有高结合亲和力的配体分子，作为参考， (ii) 将关键的蛋白质-配体结合结构纳入到分子扩散生成的引导中，采用两种有效的增强机制，即检索增强（retrieval augmentation）和自增强（self augmentation）。

在 CrossDocked 2020 数据集上的研究结果表明，IRDiff 能够生成具有更真实 3D 结构的分子，并在针对蛋白质靶点的结合亲和力方面实现最先进的性能，同时保持适当的类药性、可合成性等分子药理学属性。

二、模型介绍

设计能够与特定蛋白质靶点结合并调节其功能的配体分子，是药物发现中的一个根本性问题。针对这一问题，提出了许多新的生成方法。自回归模型通过迭代地基于目标结合位点添加原子或键来生成 3D 分子，但往往会遭受错误积累的问题，并且很难找到最优的生成顺序。为了应对自回归模型的局限性，最近的工作采用扩散模型（例如： TargetDiff），通过后处理来分配键，基于标准高斯先验对原子类型和位置的分布进行建模。这些扩散方法开发了 SE(3)-等变扩散模型，以捕捉原子之间的局部和全局空间相互作用，并且表现出比自回归模型更优异的性能，但生成的分子在满足生物学指标（如结合亲和力）等方面仍存在困难。这一困难主要来自于分子内部和分子之间的姿态的广泛搜索空间。从零开始生成分子使得生成过程更难优化，可能导致次优性能。

为了解决这些挑战，作者提出了 IRDiff （Interaction-based Retrieval-augmented Diffusion model）。IRDiff 受到机器学习中生成建模的最新进展的启发，特别是检索增强生成或上下文学习，这种方法能够通过适当的任务特定参考很好地泛化到以前未见的任务中。如下图所示，与之前仅依赖生成模型的泛化能力来处理新蛋白质靶点的方法不同，IRDiff 明确利用了一小组具有高结合亲和力的目标感知分子配体参考，以特定蛋白质为目标。通过利用参考配体与给定蛋白质之间的蛋白质-分子相互作用信息，引导扩散模型生成能紧密结合于目标口袋的配体分子。

具体来说，IRDiff 中引入了一个名为 PMINet 的蛋白质-分子相互作用网络，并用结合亲和力信号对其进行预训练，该网络通过 SE(3)-等变层和注意力层进行参数化，以捕捉蛋白质-分子对之间的相互作用信息。然后，利用预训练的 PMINet 来 (1) 检索具有高结合亲和力的蛋白质感知配体分子，作为分子参考，(2) 通过两种有效的增强机制，即检索增强和自增强，将关键的蛋白质-分子结合结构纳入到分子扩散生成中，这些生成过程是基于分子配体参考集和目标蛋白质进行的。值得注意的是，即使 PMINet 仅关注建模蛋白质与分子配体之间的序列级相互作用，IRDiff 仍然能够有效利用 PMINet 建模的蛋白质-分子相互作用进行 3D 蛋白质特异性分子生成。这种能力增强了 IRDiff 的广泛适用性和可扩展性。在 CrossDocked 2020 数据集上的大量实验表明，IRDiff 在结合相关指标方面取得了新的最先进的性能。

总的来说，IRDiff 的主要贡献是：

（1）提出了一种基于交互的检索增强 3D 分子扩散模型 IRDiff ，用于 SBDD （ structure-based drug design）任务。该模型通过有信息的外部目标感知参考引导 3D 分子生成，弥合了结合亲和力预测任务及其逆问题的鸿沟。

（2）设计了两种新颖的增强机制，即检索增强和自增强，以纳入关键的蛋白质-分子结合结构，用于目标感知的分子生成。

（3）IRDiff 能够生成紧密结合于目标口袋且具有适当分子性质的配体。在 CrossDocked 2020 数据集上的评估结果表明，模型达到了 -6.03 的平均 Vina 得分和 0.53 的平均 QED 得分，表明了结合相关指标和性质相关指标之间的显著平衡。

2.1 模型框架

IRDiff 的整体框架如下图所示。预训练的结合亲和性模型（PMINet）全面建模了蛋白质和配体之间复杂的相互作用信息，可以用作分子数据库中蛋白质感知配体参考的检索工具。作者提出了样例增强和自增强机制，通过利用配体参考和预训练的PMINet，帮助目标蛋白进行蛋白质感知的3D分子生成。

2.1.1 PMINet (protein-ligand interaction networks)

为了捕捉蛋白质和配体分子之间的相互作用，引入一个蛋白质-配体相互作用网络（即 PMINet，protein-ligand interaction networks），用于建模蛋白质-配体对的结合亲和力，该网络由 SE(3)-等变神经网络以及交叉注意力层组成。两个浅层的SE(3)-等变神经网络应用于蛋白质和配体的全连接图上，以建模分子内的相互作用。

然后使用原子的交叉注意机制，学习蛋白和小分子的分子间相互作用：

2.1.2 检索增强（retrieval-augmented）

接着，利用 PMINet 学到的基于结构的蛋白质-配体相互作用先验，识别出最合适的候选分子（即具有高结合亲和力的配体分子），从而增强后续的目标感知分子设计。更具体来说，给定特定蛋白质和一个外部分子数据库，使用 PMINet 扫描数据库，从中间检索出与给定蛋白质结合亲和力最高的一批分子作为参考（top k），可以有效促进生成配体的有效性。比如，参考分子具有较高的结合亲和力，可能揭示出一些关键信息（如：重要的氢键供体或受体），以支持强的分子内相互作用。

IRDiff 通过自增强和检索增强，充分利用了外部配体中的信息性结合先验以及蛋白质感知的交互上下文，用于生成 3D 等变分子扩散。自增强指的是通过非端到端的分子嵌入配体和蛋白质信息对生成的配体分子和目标蛋白质进行自增强。当生成过程接近完成时，此时的生成分子是和目标蛋白质具有最高结合亲和力的候选配体。检索增强策略可以提升生成过程中的参考配体质量，通过 PMINet 生成的参考分子集用于增强蛋白质和配体之间的上下文交互。通过可训练的交叉注意力机制将参考配体和生成分子的嵌入合并。

作者在蛋白质-配体复合物的 K-近邻图上应用 SE(3)-等变神经网络，以学习生成过程中的逐原子蛋白质-分子相互作用。因为蛋白质原子的坐标是已知的，所以在更新过程中保持不变。通过将蛋白质原子的质心平移为零以及 EGNN 的设计，确保逆变换的 SE(3)等变性。增强过程只会增强 SE(3)-不变的隐藏状态，而不会破坏 SE(3)-等变性。

下表中展示 IRDiff 的训练算法。模型训练需要蛋白质-配体数据集、设计好的神经网络架构、外部的分子数据库、预训练好的 PMINet 以及在每个检索池中获得的参考配体数目。首先使用 PMINet 从外部分子数据库中检索出前 k 个高结合亲和力的配体分子。在扩散时间 T 中进行采样，使蛋白质原子的质心为零。扰动初始状态一直到 t 时刻，将扩散中的 t+1步小分子和对应的蛋白质输入到 PMINet 中，得到对应的蛋白质和小分子的嵌入，将其与扩散模型中编码的对应特征融合起来。获取增强后的蛋白质和配体原子特征，与参考集中靶蛋白和小分子特征分别融合，根据融合后的特征预测原本的特征。

训练好模型之后，进行分子采样的算法如下表所示。采样时，输入给定蛋白的结合位点、外部分子库和参考池中参考配体的数目。采样配体分子的原子数目，使蛋白质原子的质心为零，初始化配体原子坐标和原子类型，将蛋白的原子类型嵌入到提取的蛋白质原子级别嵌入。进入更新过程，不断通过自增强和检索增强更新配体特征，最终获得适应给定蛋白结合位点的生成分子。

2.2 数据集、基线模型和评估指标

为了使用结合亲和力信号预训练 PMINet，使用 PDBbind v2016 数据集，该数据集经常用于结合亲和力预测任务。具体来说，3767 个复合物被选为训练集，另外 290 个复合物被选为测试集。对于分子生成，在 CrossDocked 2020 数据集上训练并评估 IRDiff 。遵循与 AR 相同的数据准备和划分方法，其中 2250 万个对接的结合复合物被精炼为高质量的对接姿势（对接姿势与真实值之间的 RMSD < 1Å），并选择了不同的蛋白质（序列相似性 < 30%）。最终获得 100,000 对用于训练的蛋白质-配体对，以及100个用于测试的蛋白质。从训练集中随机选择128 个配体用于检索，并为每个蛋白质选择预测结合亲和力排名第一的配体作为参考。

作者选择了当前先进的基于结构的分子生成基线模型和 IRDiff 进行比较。LiGAN 是一种条件变分自编码器（VAE）模型，基于蛋白质-配体结构的原子密度网格表示进行训练。AR 和 Pocket2Mol 是自回归方案，生成的 3D 分子原子是基于蛋白质口袋和之前生成的原子。TargetDiff 和 DecomposeDiff 是最近的最先进的非自回归扩散基础 SBDD 模型。

对生成分子的评估主要从三个角度：分子结构、靶标结合亲和力和分子属性。在分子结构方面，计算生成分子和测试集中提供的真实分子的原子/键距离的经验分布之间的 Jensen-Shannon 散度（JSD）。为估算靶标结合亲和力，采用 AutoDock Vina 计算并报告结合相关指标的均值和中值，包括 Vina 评分（Vina Score）、Vina 最小值（Vina Min）、Vina 重对接（Vina Dock）和高亲和力（High Affinity）。Vina 评分直接基于生成的 3D 分子估计结合亲和力；Vina 最小值在估计前执行局部结构最小化；Vina 重对接涉及一个额外的重对接过程，反映了最佳的结合亲和力；高亲和力衡量每个测试蛋白中生成的分子比真实分子结合更好的比例。为了评估分子属性，采用 QED、SA 和多样性（Diversity）作为指标。QED 是一个简单的药物相似性定量估计，结合了几种理想的分子属性；SA 衡量了合成配体的难度；多样性则通过计算所有生成配体之间的平均成对差异来评估。

2.3 模型性能

作者比较了 IRDiff 与代表性方法在分子结构方面的表现。下图展示了生成分子的全原子成对距离分布。IRDiff 在生成分子的全原子成对距离分布中与真实分子的 JSD 为 0.08，优于两个强基线方法 Pocket2Mol 和 TargetDiff，表明其有效捕捉到了真实的原子距离。

接着，作者计算了生成分子的不同键的分布，并与对应的真实分子的经验分布进行了比较，结果如下表所示，-、= 和：分别表示单键、双键和芳香键。IRDiff 在大多数主要键类型上与 DecompDiff 表现相当，并显著优于其他基线，表明 IRDiff 在生成稳定的分子结构方面具有巨大潜力。将这一结果归因于增强机制，它为分子生成提供了真实的 3D 配体模板。

作者评估了 IRDiff 在结合亲和力方面的有效性。从下表可以看出，IRDiff 在结合相关指标上优于基线方法。特别是，IRDiff 在平均和中位 Vina 评分上，分别以 17.3% 和 46.6% 的显著优势超过了强自回归方法 Pocket2Mol，并在平均和中位 Vina 评分上分别以 6.3% 和 14.1% 超过了强扩散方法 DecompDiff。在高亲和力方面，发现平均而言，67.4% 的 IRDiff 生成分子在测试集中表现出比真实分子更好的结合亲和力，显著优于其他基线方法。这些提升表明，IRDiff 能够有效利用外部的蛋白质-配体相互作用，生成具有更高靶标结合亲和力的分子。从评估结果中可以看到在基线模型方法中，分子属性相关指标（QED 和 SA）与结合相关指标之间存在权衡。TargetDiff 和 DecompDiff 在结合相关指标上表现优于 AR 和 Pocket2Mol，但在 QED 和 SA 得分上落后。相比之下，IRDiff 不仅达到了最先进的结合相关分数，还保持了适当的 QED 分数，实现了比 TargetDiff 和 DecompDiff 更好的平衡。然而，IRDiff 对 QED 和 SA 的重视较少，因为它们通常在实际药物发现过程中仅作为粗略的筛选指标，只要它们在合理范围内即可。

下图展示了一些生成分子的示例及其属性。IRDiff 生成的分子结构有效，并且与目标蛋白具有合理的结合位点，是有前景的候选配体。

作者研究了 IRDiff 的自增强和检索增强机制的影响。下表展示了自增强和检索增强在 IRDiff 中的有效性。移除了 IRDiff 的增强机制的模型作为基线进行比较。结果表明，仅仅应用增强机制而不使用预训练的 PMINet 会降低生成性能，因为它不包含用于自我优化的有用蛋白质-配体相互作用信息。相反，自增强机制由于预训练 PMINet 带来的相互作用知识，显著提升了结合相关和属性相关的指标。此外，检索增强也相较于基线有显著改进，表明外部的目标感知参考确实提供了合适的参照，并促进了分子生成的优化。检索增强对属性相关指标帮助不大，因为主要关注示例中的有用结合结构。结合使用检索增强和自增强机制达到了最佳的结合相关指标，证明了 IRDiff 中这两种互补增强机制的有效性。

下图展示了生成分子与数据库中的所有分子和对应的参考分子的分子指纹相似度。提供了生成配体与其相应参考之间相似性的定量分析，生成分子和参考分子的相似性高于数据库分子的，表明这些参考分子确实作为示例引导了分子生成。

下图展示了生成分子和相应的参考分子，两者间相同的关键子片段用相同颜色标注出来。这表明模型可能会自动选择参考配体中对靶标具有高结合亲和力的关键子结构，并将其重新组合在适当位置，以生成新的分子。

作者通过对变量 n 和 k 的消融实验研究了 IRDiff 中 3D 分子数据库对检索的影响，其中 n 表示分子数据库的大小，k 表示参考样本的数量。从下表中的结果可以看出，较大的 n 有助于提升 IRDiff 在结合相关指标上的表现，因为更高的多样性允许更多的结合相关线索（子结构）用于增强生成过程。然而，简单增加 k 并没有明显的提升，因为使用更多的分子参考可能会引入更多的噪声，干扰生成过程。

作者还研究了参考配体的结合亲和力对 IRDiff 生成性能的影响。在实验中，选择了具有最高（即 k=1）目标结合亲和力的配体作为参考，设置参考池的大小 n 为128。在这里，参考配体被替换为在 PMINet 提供的排名列表中具有最低结合亲和力的配体。这些实验是在没有使用自增强机制的 IRDiff 上进行的。正如下表所示，使用低亲和力参考（PMINet 预测的参考分子与相应蛋白质口袋之间的平均结合亲和力为 3.95）的分子数据库在结合相关指标上的表现比使用高亲和力参考（平均结合亲和力为 7.24）要差。这证明了在 IRDiff 中基于相互作用的检索设计和参考利用的有效性。

3.1 安装环境

复制代码项目：

git clone https://github.com/YangLing0818/IRDiff.git

项目提供了项目运行的环境配置文件 irdiff.yml，在配置项目运行的虚拟环境时，通常使用项目的名字作为虚拟环境的名字，开头用大写字母，所以我们将名字从原本的 ipdiff 手动修改为 IPDiff。配置文件中部分内容展示如下：

name: IRDiff
channels:
  - pytorch
  - pyg
  - conda-forge
  - defaults
dependencies:
  - _ipython_minor_entry_point=8.7.0=h3b92ee0_0
  - _libgcc_mutex=0.1=conda_forge
  - _openmp_mutex=4.5=2_kmp_llvm

使用提供的 irdiff.yml 创建 IRDiff 环境，命令如下：

conda env create -f irdiff.yml

但是安装过程中的依赖库版本有冲突，报错如下：

LibMambaUnsatisfiableError: Encountered problems while solving:
  - package cryptography-37.0.1-py38h9ce1e76_0 requires openssl <1.1.2a, but none of the providers can be installed

Could not solve for environment specs
The following packages are incompatible
├─ cryptography ==37.0.1 py38h9ce1e76_0 is installable and it requires
│  └─ openssl <1.1.2a , which can be installed;
└─ openssl ==3.1.0 hd590300_2 is not installable because it conflicts with any installable versions previously reported.

irdiff.yml 中 cryptography ==37.0.1 指定的版本需要 openssl <1.1.2a，配置文件中指定了 openssl ==3.1.0，有版本冲突。把 irdiff.yml 中第 42 行中，cryptography 指定的版本信息删除，让 conda 自动寻找匹配的版本.

cryptography=37.0.1=py38h9ce1e76_0

改为

cryptography

修改之后，继续安装环境

conda env create -f irdiff.yml

但仍报错 autodocktools 的安装错误，如下

Collecting package metadata (repodata.json): done
Solving environment: done


Downloading and Extracting Packages:


Preparing transaction: done
Verifying transaction: done
Executing transaction: - By downloading and using the CUDA Toolkit conda packages, you accept the terms and conditions of the CUDA End User License Agreement (EULA): https://docs.nvidia.com/cuda/eula/index.html


/ 


    Installed package of scikit-learn can be accelerated using scikit-learn-intelex.
    More details are available here: https://intel.github.io/scikit-learn-intelex


    For example:


        $ conda install scikit-learn-intelex
        $ python -m sklearnex my_application.py


    


done
Installing pip dependencies: \ Ran pip subprocess with arguments:
['/workspace/anaconda3/envs/IRDiff/bin/python', '-m', 'pip', 'install', '-U', '-r', '/workspace/xxxx/projects/IRDiff/condaenv.hodd02ab.requirements.txt', '--exists-action=b']
Pip subprocess output:
Looking in indexes: https://mirrors.cloud.tencent.com/pypi/simple


Pip subprocess error:
ERROR: Could not find a version that satisfies the requirement autodocktools-py3==0+unknown (from versions: none)
ERROR: No matching distribution found for autodocktools-py3==0+unknown


failed


CondaEnvException: Pip failed

通过 pip 安装的 autodocktools-py3 无法找到合适的版本，进入虚拟环境手动安装

conda activate IRDiff
python -m pip install git+https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3

autodocktools-py3 安装成功，打印下面内容：

Looking in indexes: https://mirrors.cloud.tencent.com/pypi/simple
Collecting git+https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3
  Cloning https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3 to /tmp/pip-req-build-3mmsr8we
  Running command git clone -q https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3 /tmp/pip-req-build-3mmsr8we
  Resolved https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3 to commit a62a5d98116f1590183b58a9ad732b997cf2579c
Building wheels for collected packages: AutoDockTools-py3
  Building wheel for AutoDockTools-py3 (setup.py) ... done
  Created wheel for AutoDockTools-py3: filename=AutoDockTools_py3-1.5.7.post1+12.ga62a5d9-py3-none-any.whl size=984166 sha256=9b166e832e3641398aa7e4612877719b44aec816a91e12aec934d2aea9006883
  Stored in directory: /tmp/pip-ephem-wheel-cache-il3coo94/wheels/88/95/bf/1ff0c68c3d3146de99a9b93ad95be4738164f5980b5a415661
Successfully built AutoDockTools-py3
Installing collected packages: AutoDockTools-py3
Successfully installed AutoDockTools-py3-1.5.7.post1+12.ga62a5d9

剩余的依赖库，也通过下面命令安装：

pip install meeko==0.1.dev3 pdb2pqr==3.6.1 vina==1.2.2 propka==3.5.0 mmcif-pdbx==2.0.1

安装过程输出如下：

Looking in indexes: https://mirrors.cloud.tencent.com/pypi/simple
Collecting meeko==0.1.dev3
  Using cached https://mirrors.cloud.tencent.com/pypi/packages/89/84/827a0ea853ba0970f0287eb09ba7d07186590a510d10a8b63a2d884ac447/meeko-0.1.dev3-py2.py3-none-any.whl (45 kB)
Collecting pdb2pqr==3.6.1
  Using cached https://mirrors.cloud.tencent.com/pypi/packages/3a/7c/3bbf1f414f70cbb14b1ee4da74061ae5a0323cd4cb037a76642fa23d2e2f/pdb2pqr-3.6.1-py2.py3-none-any.whl (208 kB)
Collecting vina==1.2.2
  Using cached https://mirrors.cloud.tencent.com/pypi/packages/18/38/d197002b15b4190d005da557c00dd96af320c0fb49c5b1a9130d01a90e1c/vina-1.2.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.0 MB)
Collecting propka==3.5.0
  Using cached https://mirrors.cloud.tencent.com/pypi/packages/c4/e2/5c096dc02874a217b26d01cecb701278c5fbf847e473f4870a711563eb87/propka-3.5.0-py3-none-any.whl (98 kB)
Collecting mmcif-pdbx==2.0.1
  Using cached https://mirrors.cloud.tencent.com/pypi/packages/03/8e/ff50191d2210faac7df6b2e66cf39e021bd7051bd9389b341b1d560cae49/mmcif_pdbx-2.0.1-py2.py3-none-any.whl (20 kB)
Requirement already satisfied: numpy>=1.18 in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from meeko==0.1.dev3) (1.22.3)
Requirement already satisfied: requests in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from pdb2pqr==3.6.1) (2.27.1)
Requirement already satisfied: docutils<0.18 in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from pdb2pqr==3.6.1) (0.17.1)
Requirement already satisfied: charset-normalizer~=2.0.0 in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from requests->pdb2pqr==3.6.1) (2.0.4)
Requirement already satisfied: idna<4,>=2.5 in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from requests->pdb2pqr==3.6.1) (3.3)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from requests->pdb2pqr==3.6.1) (1.26.15)
Requirement already satisfied: certifi>=2017.4.17 in /workspace/anaconda3/envs/IRDiff/lib/python3.8/site-packages (from requests->pdb2pqr==3.6.1) (2022.12.7)
Installing collected packages: propka, mmcif-pdbx, vina, pdb2pqr, meeko
Successfully installed meeko-0.1.dev3 mmcif-pdbx-2.0.1 pdb2pqr-3.6.1 propka-3.5.0 vina-1.2.2

注：安装过程要注意 pyg 的版本，使用 2.0.4 。其他版本可能会在训练时出现问题。

3.2 分子生成案例测试

IRDiff 的数据预处理和数据集和 TargetDiff 一致，相关文件也保存在谷歌网盘中，链接为

https://drive.google.com/drive/folders/1j21cc7-97TedKh_El5E34yI8o5ckI7eK?usp=share_link ，内容具体如下图所示。下载好的 crossdocked_pocket10_pose_split.pt 和

crossdocked_v1.1_rmsd1.0_pocket10_processed_final.lmdb 放在 ./data 文件夹中。

项目下载到本地，路径设置有所改变，我们把所有的路径设置到适应当前项目的相对路径，具体如下：

（1）cal_metrics_from_pt.py 中第 8 行的

eval_path = '/path/to/eval_results/'

修改为

eval_path = './eval_results/'

（2）./configs/training.yml 中第 4-5 行的

path: /path/to/crossdocked_v1.1_rmsd1.0
split: /path/to/crossdocked_pocket10_pose_split.pt

修改为

path:  ./data/crossdocked_v1.1_rmsd1.0
split: ./data/crossdocked_pocket10_pose_split.pt

3.2.1 内置案例

测试集中 index 为 0 的体系（PDB : 2Z3H，即 ./datasets/test_set/BSD_ASPTE_1_130_0 的测试体系）的口袋以及原来小分子的展示如下图，我们把这个蛋白作为测试分子生成的内置案例。

口袋中配体的 2D 结构如下：

分子采样的脚本是 sample_split.py，使用编号为 0 的蛋白进行分子生成，命令如下：

python sample_split.py \
  --start_index 0 \
  --end_index 0 \
  --batch_size 25 \
  --test_prompt_indices_path ./src/test_prompt_ligand_indices_top3.pt \
  --result_path ./results/result_0

--start_index 0 和 --end_index 0 指定分子生成的起始和结束编号，这里指定编号为 0 的蛋白，--result_path ./results/result_0 指定生成分子保存位置。分子生成命令默认配置文件为 ./configs/sampling.yml，具体配置如下：

model:
  checkpoint: ./pretrained_models/pretrained-IRDiff.pt


sample:
  seed: 2021
  num_samples: 100
  num_steps: 1000
  pos_only: False
  center_pos_mode: protein
  sample_num_atoms: prior

分子采样命令报错如下：

Traceback (most recent call last):
  File "sample_split.py", line 17, in <module>
    from models.molopt_score_model import ScorePosNet3D, log_sample_categorical
  File "/workspace/xxxx/projects/IRDiff/models/molopt_score_model.py", line 11, in <module>
    from models.attn import CrossAttention, RetAugmentationLinearAttention
  File "/workspace/xxxx/projects/IRDiff/models/attn.py", line 4, in <module>
    from einops import rearrange, repeat
ModuleNotFoundError: No module named 'einops'

根据报错提示，环境中缺少 'einops'，手动安装命令如下：

conda install einops

内置案例采样命令正常执行，batch_size 设置为 25 ，显存占用约为 19 GB，采样 100 个分子，花费大约 50 分钟，采样输出如下：

[2024-09-14 05:03:25,154::sampling::INFO] {'model': {'checkpoint': './pretrained_models/pretrained-IRDiff.pt'}, 'sample': {'seed': 2021, 'num_samples': 100, 'num_steps': 1000, 'pos_only': False, 'center_pos_mode': 'protein', 'sample_num_atoms': 'prior'}}
[2024-09-14 05:03:28,726::sampling::INFO] Training Config: {'data': {'topk_prompt': 3, 'name': 'pl', 'path': '/workspace/xxxx/projects/IRDiff/data/crossdocked_v1.1_rmsd1.0', 'split': '/workspace/xxxx/projects/IRDiff/data/crossdocked_pocket10_pose_split.pt', 'train_prompt_path': './src/train_prompt_ligand_indices_top3.pt', 'test_prompt_path': './src/test_prompt_ligand_indices_top3.pt', 'val_prompt_path': './src/val_prompt_ligand_indices_top3.pt', 'transform': {'ligand_atom_mode': 'add_aromatic', 'random_rot': False}}, 'net_cond': {'ckpt_path': './pretrained_models/PMINet', 'hidden_dim': 128}, 'model': {'cond_dim': 128, 'topk_prompt': 3, 'model_mean_type': 'C0', 'beta_schedule': 'sigmoid', 'beta_start': 1e-07, 'beta_end': 0.002, 'v_beta_schedule': 'cosine', 'v_beta_s': 0.01, 'num_diffusion_timesteps': 1000, 'loss_v_weight': 100.0, 'sample_time_method': 'symmetric', 'time_emb_dim': 0, 'time_emb_mode': 'simple', 'center_pos_mode': 'protein', 'node_indicator': True, 'model_type': 'uni_o2', 'num_blocks': 1, 'num_layers': 9, 'hidden_dim': 128, 'n_heads': 16, 'edge_feat_dim': 4, 'num_r_gaussian': 20, 'knn': 32, 'num_node_types': 8, 'act_fn': 'relu', 'norm': True, 'cutoff_mode': 'knn', 'ew_net_type': 'global', 'num_x2h': 1, 'num_h2x': 1, 'r_max': 10.0, 'x2h_out_fc': False, 'sync_twoup': False}, 'train': {'seed': 2021, 'batch_size': 4, 'val_batch_size': 1, 'num_workers': 4, 'n_acc_batch': 1, 'max_iters': 1000000, 'val_freq': 5000, 'pos_noise_std': 0.1, 'max_grad_norm': 8.0, 'bond_loss_weight': 1.0, 'optimizer': {'type': 'adam', 'lr': 0.0005, 'weight_decay': 0, 'beta1': 0.95, 'beta2': 0.999}, 'scheduler': {'type': 'plateau', 'factor': 0.6, 'patience': 10, 'min_lr': 1e-06}}}
[2024-09-14 05:03:28,737::sampling::INFO] Successfully load the dataset (size: 100)!
Restored from ./pretrained_models/PMINet with 0 missing and 0 unexpected keys
[2024-09-14 05:03:28,829::sampling::INFO] Successfully load the model! ./pretrained_models/pretrained-IRDiff.pt
sampling: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [12:35<00:00,  1.32it/s]
sampling: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [12:33<00:00,  1.33it/s]
sampling: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [12:30<00:00,  1.33it/s]
sampling: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [12:32<00:00,  1.33it/s]
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [50:36<00:00, 759.22s/it]
[2024-09-14 05:54:05,761::sampling::INFO] Sample done!

在 IRDiff 生成和评价分子的过程中，并没有保存生成的分子构象，所以我们修改评价分子的脚本 eval_split.py，加入保存分子构象的代码。

原代码：

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    ...
    if args.save:
        torch.save({
            'stability': validity_dict,
            'bond_length': all_bond_dist,
            'all_results': results
        }, os.path.join(result_path, f'metrics_{args.eval_step}_{eval_start_index}-to-{eval_end_index}.pt'))

在脚本结尾添加保存分子构象的代码，用以保存生成的分子构象。保存的分子文件名为：Generated_molecules.sdf，路径为 ./eval_results/Generated_molecules.sdf。修改后的代码：

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    ...
    if args.save:
        torch.save({
            'stability': validity_dict,
            'bond_length': all_bond_dist,
            'all_results': results
        }, os.path.join(result_path, f'metrics_{args.eval_step}_{eval_start_index}-to-{eval_end_index}.pt'))


    # 保存生成的分子
    mol_sdfs = []
    for res in results:
        # 保存生成分子的构象
        mol = res['mol']
        if args.docking_mode == 'vina_score':
          # 给分子添加 vina_score 属性
          vina_score = res['vina']['score_only'][0]['affinity']
          mol.SetProp('vina_score', str(vina_score))
        mol_sdfs.append(mol)
    if args.docking_mode == 'vina_score':
      # 每个 mol 对象都有 'vina_score' 属性，按 vina_score 降序排列
      mol_sdfs = sorted(mol_sdfs, key=lambda mol: float(mol.GetProp('vina_score')), reverse=False)
    
    # 把生成分子保存成 sdf 格式
    w = Chem.SDWriter(f'./{result_path}/Generated_molecules.sdf')
    for mol in mol_sdfs:
        w.write(mol)
    w.close()
    print('Generated molecules saved as sdf format!')

接着，使用 qvina 对接打分评价生成分子，--sample_path ./results/result_0 指定生成分子的路径，--protein_root ./data/test_set 指定生成分子需要对接的测试蛋白位置，

--docking_mode qvina 指定使用 qvina 对接生成分子。具体命令如下：

python eval_split.py \
  --eval_start_index 0 \
  --eval_end_index 0 \
  --sample_path ./results/result_0 \
  --docking_mode qvina \
  --protein_root ./data/test_set \
  --exhaustiveness 90

分子评价返回结果如下：

[2024-09-14 09:05:56,052::evaluate::INFO] Load generated data done! sample_id[0:0] examples for evaluation.
[2024-09-14 09:33:41,158::evaluate::INFO] Evaluate done! 100 samples in total.
[2024-09-14 09:33:41,158::evaluate::INFO] mol_stable:	0.5200
[2024-09-14 09:33:41,158::evaluate::INFO] atm_stable:	0.9629
[2024-09-14 09:33:41,158::evaluate::INFO] recon_success:	1.0000
[2024-09-14 09:33:41,158::evaluate::INFO] eval_success:	0.9300
[2024-09-14 09:33:41,158::evaluate::INFO] complete:	0.9500
[2024-09-14 09:33:41,162::evaluate::INFO] JS bond distances of complete mols: 
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-6|4:	0.4319
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-6|1:	0.3086
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-8|1:	0.4796
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-7|1:	0.3287
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-8|2:	0.4293
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-6|2:	0.4114
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-7|4:	0.7288
[2024-09-14 09:33:41,162::evaluate::INFO] JSD_6-7|2:	0.6420
[2024-09-14 09:33:41,186::evaluate::INFO] JSD_CC_2A:	0.3409
[2024-09-14 09:33:41,187::evaluate::INFO] JSD_All_12A:	0.1006
[2024-09-14 09:33:41,187::evaluate::INFO] Atom type JS: 0.1543
[2024-09-14 09:33:41,360::evaluate::INFO] Number of reconstructed mols: 100, complete mols: 95, evaluated mols: 93
[2024-09-14 09:33:41,360::evaluate::INFO] QED:   Mean: 0.480 Median: 0.486
[2024-09-14 09:33:41,360::evaluate::INFO] SA:    Mean: 0.527 Median: 0.540
[2024-09-14 09:33:41,361::evaluate::INFO] Vina:  Mean: -2.254 Median: -4.300
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 3 ratio: 0.000
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 4 ratio: 0.043
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 5 ratio: 0.602
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 6 ratio: 0.849
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 7 ratio: 0.591
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 8 ratio: 0.204
[2024-09-14 09:33:41,361::evaluate::INFO] ring size: 9 ratio: 0.065

从返回结果可以看出，能够重构的生成分子有 93 个，QED 和 SA 均值分别为 0.480 和 0.527，生成分子中包含五、六、七环的比例分别为 0.602、0.849 和 0.591。评估结果和生成的分子构象默认保存在 ./eval_results/ 中。因为所有评价结果以及生成分子会输出保存到 ./eval_results/，为防止后面的案例覆盖结果，我们把 ./eval_results 重命名为 ./eval_results_0。

所有生成的分子如下：

IRDiff_2z3h_outputs

qvina score 排名前 3 的分子的 2D 结构如下，对应的 qvina score 分数分别为 -9.0， -8.7， -8.4：

qvina score 排名前 3 的分子在口袋中的 Pose 如下，蓝色的是参考分子，紫红色的是生成分子，生成分子主要占据口袋内的空间，口袋的开放空间没有参考分子占据的多：

3.2.2 自定义的测试案例

我们选择 3WZE 作为自己的测试案例，使用 PyMol 把 3WZE 的配体周围 10 Å 的范围作为口袋，保存为 ./3wze/pocket_3wze.pdb，3WZE 口袋与原来配体小分子的如下图：

口袋中分子的 2D 结构，如下：

项目提供 scripts/sample_for_pocket.py 可以基于蛋白口袋对分子采样


python scripts/sample_for_pocket.py \
    configs/sampling.yml \
    --pdb_path ./3wze/pocket_3wze.pdb \
    --batch_size 25 \
    --result_path ./results/result_3wze

自定义案例分子采样命令报错如下：

[2024-09-15 10:10:43,784::evaluate::INFO] {'model': {'checkpoint': './pretrained_models/pretrained-IRDiff.pt'}, 'sample': {'seed': 2021, 'num_samples': 100, 'num_steps': 1000, 'pos_only': False, 'center_pos_mode': 'protein', 'sample_num_atoms': 'prior'}}
Traceback (most recent call last):
  File "scripts/sample_for_pocket.py", line 51, in <module>
    logger.info(f"Training Config: {ckpt['config']}")
KeyError: 'config'

经过仔细检查发现，scripts/sample_for_pocket.py 文件中的代码并不是 IRDiff 的基于口袋进行分子生成的代码，仍然是原来 TargetDiff 的代码。这部分代码，也不能进行 prompt 均为 None 的情况下进行分子生成，因为，由于模型 ScorePosNet3D 作者没有兼容 prompt 均为 None 的情况。因此，需要对整个 scripts/sample_for_pocket.py 文件进行改写，并重命名为 scripts/sample_for_pocket_wufeil.py 文件

经过修改代码以后，使用的生成命令如下：

python scripts/sample_for_pocket_wufeil.py \
    configs/sampling.yml \
    --pdb_path ./3wze/pocket_3wze.pdb \
    --reference_mol_path ./3wze/3wze-mol.sdf \
    --batch_size 25 \
    --result_path ./results/result_3wze

正常运行生成的分子保存在 ./results/result_3wze/sample_0.pt 文件内。

运行以下命令对生成的分子进行评估（注：evaluate_for_pocket.py 文件来源于IPDiff 项目）：

python ./scripts/evaluate_for_pocket.py \
  ./results/result_3wze/sample_0.pt \
  --verbose True \
  --protein_path ./3wze/pocket_3wze.pdb \
  --docking_mode qvina \
  --exhaustiveness 16

运行输出：

[2024-09-27 08:12:36,538::evaluate::INFO] Evaluate done! 100 samples in total.
[2024-09-27 08:12:36,538::evaluate::INFO] mol_stable:   0.3700
[2024-09-27 08:12:36,538::evaluate::INFO] atm_stable:   0.9313
[2024-09-27 08:12:36,538::evaluate::INFO] recon_success:        1.0000
[2024-09-27 08:12:36,538::evaluate::INFO] eval_success: 0.8600
[2024-09-27 08:12:36,538::evaluate::INFO] complete:     0.8700
[2024-09-27 08:12:36,541::evaluate::INFO] JS bond distances of complete mols: 
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-6|4:    0.3019
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-6|1:    0.4614
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-8|1:    0.4043
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-7|1:    0.3785
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-8|2:    0.4333
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-6|2:    0.2926
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-7|4:    0.6505
[2024-09-27 08:12:36,541::evaluate::INFO] JSD_6-7|2:    0.4496
[2024-09-27 08:12:36,553::evaluate::INFO] JSD_CC_2A:    0.3147
[2024-09-27 08:12:36,553::evaluate::INFO] JSD_All_12A:  0.0923
[2024-09-27 08:12:36,553::evaluate::INFO] Atom type JS: 0.1842
[2024-09-27 08:12:36,724::evaluate::INFO] Number of reconstructed mols: 100, complete mols: 87, evaluated mols: 86
[2024-09-27 08:12:36,724::evaluate::INFO] QED:   Mean: 0.609 Median: 0.631
[2024-09-27 08:12:36,724::evaluate::INFO] SA:    Mean: 0.654 Median: 0.640
[2024-09-27 08:12:36,724::evaluate::INFO] Vina:  Mean: -6.256 Median: -7.600
[2024-09-27 08:12:36,724::evaluate::INFO] ring size: 3 ratio: 0.000
[2024-09-27 08:12:36,725::evaluate::INFO] ring size: 4 ratio: 0.047
[2024-09-27 08:12:36,725::evaluate::INFO] ring size: 5 ratio: 0.477
[2024-09-27 08:12:36,725::evaluate::INFO] ring size: 6 ratio: 0.872
[2024-09-27 08:12:36,725::evaluate::INFO] ring size: 7 ratio: 0.453
[2024-09-27 08:12:36,725::evaluate::INFO] ring size: 8 ratio: 0.163
[2024-09-27 08:12:36,725::evaluate::INFO] ring size: 9 ratio: 0.035
Generated molecules saved as sdf format in ./results/result_3wze/eval_results/Generated_molecules.sdf!

从返回结果可以看出，能够重构的生成分子有 86 个，QED 和 SA 均值分别为 0.609 和 0.654，生成分子中包含五、六、七环的比例分别为 0.477、0.872 和 0.453。评估结果和生成的分子构象默认保存在 ./results/result_3wze/eval_results 文件夹内。

所有生成的分子如下：

IRDiff_3wze_outputs

qvina score 排名前 3 的分子的 2D 结构如下，对应的 qvina score 分数分别为 -11.7， -11.1， -10.8：

qvina score 排名前 3 的分子在口袋中的 Pose 如下，蓝色的是参考分子，绿色的是生成分子：

四、总结

在这项工作中，作者首次提出了一种基于交互的检索增强 3D 分子扩散模型 IRDiff ，用于基于结构的药物设计（SBDD）。IRDiff 利用目标感知的参考配体来增强 3D 分子扩散生成，通过有效的自我增强和检索增强机制，显著提高了由 Vina 测量的结合亲和力，同时保持了适当的分子属性。在未来的工作中，可以把其他与结合相关的信息纳入生成过程中。

我们实际的测试表明，IRDiff 生成的分子在 qvina 打分上，不如之前 DeepICL 等模型。生成的分子类药性比较差，存在大量不正确的键、键角。在相互作用方面，溶剂区的片段模式也没有得到保持。只能说，IRDiff 属于概念验证模型，验证了这么一种方法的科学可执行性。

完整的测评文档及此项目可运行的代码详见附件资源。