foldingdiff - 用于生成蛋白质主链结构的扩散模型

我们提出了一种用于生成新型蛋白质主链结构的扩散模型。更多详情请参阅我们在 arXiv 上发布的预印本。我们还在 HuggingFace spaces 和 SuperBio 上提供了经过训练的模型版本，您只需使用浏览器即可开始生成蛋白质结构！

安装

要安装，请使用 git clone 克隆本仓库。该软件使用 Python 编写，主要依赖 PyTorch、PyTorch Lightning 和 HuggingFace 的 transformers 库。所需的 Conda 环境已在 environment.yml 文件中定义。请确保已安装 Conda（或 mamba），克隆本仓库后运行以下命令：

conda env create -f environment.yml
conda activate foldingdiff
pip install -e ./  # 确保当前目录包含 setup.py

下载数据

由于文件体积较大，部分数据未打包在 Git 仓库中。这些数据并非采样所必需（只要不使用 --testcomparison 选项，见下文）；它们仅用于训练您自己的模型。我们在 data 目录下提供了一个脚本，用于下载所需的 CATH 数据。

# 下载 CATH 数据集
cd data  # 确保位于代码库中的 data 子目录
chmod +x download_cath.sh
./download_cath.sh

如果 .sh 脚本中的下载链接无法使用，该 tarball 也镜像在以下 Dropbox 链接中：Dropbox 链接。

训练模型

要使用 CATH 数据集训练您自己的模型，请结合 bin/train.py 脚本与 config_jsons 目录下的某个 JSON 配置文件（或自行编写）。示例如下：

python bin/train.py config_jsons/cath_full_angles_cosine.json --dryrun

默认情况下，训练脚本会在开始训练前计算每个时间步的 KL 散度，这在时间步较多时会非常耗时。若跳过此步骤，可添加 --dryrun 标志。模型输出将保存在 results 文件夹中，主要文件包括：

results/
    - config.json           # 包含 HuggingFace BERT 模型本身的配置文件
    - logs/                 # 包含训练日志
    - models/               # 包含模型检查点。默认存储验证损失和训练损失最好的各 5 个模型
    - training_args.json    # 完整的运行参数，可用于复现实验

预训练模型

我们提供了在 CATH 数据集上训练的模型权重。这些权重存储在 HuggingFace 模型中心的 wukevin/foldingdiff_cath 上。以下代码片段展示了如何加载该模型、加载数据（假设已下载）并执行前向传播：

from huggingface_hub import snapshot_download
from torch.utils.data.dataloader import DataLoader
from foldingdiff import modelling
from foldingdiff import datasets as dsets

# 加载模型（文件将被缓存以供后续调用）
m = modelling.BertForDiffusion.from_dir(snapshot_download("wukevin/foldingdiff_cath"))

# 加载数据集
# 在加载过程中，我们会尝试并行计算内部角度。这可能会抛出类似以下警告；这是正常现象。
# WARNING:root:Illegal values for omega in /home/*/projects/foldingdiff-main/data/cath/dompdb/2ebqA00 -- skipping
# 计算一次后，结果会保存为 .pkl 文件，存放在 foldingdiff 源代码目录下，以便下次更快加载。
clean_dset = dsets.CathCanonicalAnglesOnlyDataset(pad=128, trim_strategy='randomcrop')
noised_dset = dsets.NoisedAnglesDataset(clean_dset, timesteps=1000, beta_schedule='cosine')
dl = DataLoader(noised_dset, batch_size=32, shuffle=False)
x = iter(dl).next()

# 前向传播
predicted_noise = m(x['corrupted'], x['t'], x['attn_mask'])

采样蛋白质主链

要采样蛋白质主链，请使用 bin/sample.py 脚本。以下是一些使用上述预训练权重进行采样的示例命令。

# 采样长度范围为 [50, 128)，每段长度采样 10 条序列，批次大小为 512 —— 可重现论文中的结果
python ~/projects/foldingdiff/bin/sample.py -l 50 128 -n 10 -b 512 --device cuda:0

此命令将运行托管在 wukevin/foldingdiff_cath 上的训练模型，并生成不同长度的序列。如果您希望加载测试数据集并在生成的图表中包含测试链，可使用 --testcomparison 选项；请注意，这需要先下载 CATH 数据集，详见上文。运行 sample.py 后，将在其所在目录下创建如下文件结构：

some_dir/
    - plots/            # 包含训练/生成角度分布对比图
    - sampled_angles/   # 包含采样角度的 .csv.gz 文件
    - sampled_pdb/      # 包含将采样角度转换为笛卡尔坐标后的 .pdb 文件
    - model_snapshot/   # 包含用于生成结果的模型副本

若未指定 --device，则默认使用第一个 GPU 设备 cuda:0；使用 --device cpu 可在 CPU 上运行（但速度会非常慢）。以下表格显示了我们机器上的运行时间估算。

设备	采样 512 条序列的预计时间
Nvidia RTX 2080Ti	7 分钟
i9-9960X (16 核物理核心)	2 小时

最大训练相似度 TM 分数

生成序列后，我们可以计算 TM 分数来评估生成序列与原始序列的相似性。这通过 bin/tmscore_training.py 脚本完成，且需要提前下载数据（见上文）。

可视化扩散“折叠”过程

上述采样代码也可以使用 --fullhistory 标志运行，从而在每个 sampled_angles 和 sampled_pdb 文件夹下生成一个名为 sample_history 的附加子目录，其中包含与采样过程中每个时间步相对应的 pdb 和 csv 文件。例如，这些 pdb 文件可以传递给 foldingdiff/pymol_vis.py 脚本，以生成折叠过程的 GIF 动画（如上所示）。执行此操作的示例命令如下：

python ~/projects/foldingdiff/foldingdiff/pymol_vis.py pdb2gif -i sampled_pdb/sample_history/generated_0/*.pdb -o generated_0.gif

注意：该脚本与其他绘图代码分开存放，因为它依赖于 PyMOL；您可以自行安装或激活 PyMOL，也可以使用 PyMOL 开源版 设置环境。

评估生成骨架的设计性

评估生成骨架质量的一种方法是通过其“设计性”。这指的是我们是否能够设计出一条能够折叠成所设计骨架的氨基酸链。为了评估这一点，我们使用反向折叠模型生成预测会折叠成我们生成骨架的氨基酸序列，并检查这些生成的序列是否确实能够折叠成与我们的骨架结构相当的构象。

反向折叠

反向折叠是指预测能够产生给定蛋白质骨架结构的氨基酸序列的任务。我们在这一步骤中评估了两种不同的方法：ProteinMPNN 和 ESM-IF1；我们发现 ProteinMPNN 的性能显著更高。在我们的分析中，我们为 FoldingDiff 生成的每个结构生成 8 条不同的氨基酸序列。

ESM-IF1

我们为 ESM-IF1 使用了一个独立的 conda 环境；有关设置细节，请参阅此 Jupyter Notebook。我们发现以下一系列命令在我们的机器上有效：

mamba create -n inverse python=3.9 pytorch cudatoolkit pyg -c pytorch -c conda-forge -c pyg
conda activate inverse
mamba install -c conda-forge biotite
pip install git+https://github.com/facebookresearch/esm.git

此后，我们 cd 到包含上一步创建的 sampled_pdb 目录的文件夹，并运行：

python ~/projects/foldingdiff/bin/pdb_to_residues_esm.py sampled_pdb -o esm_residues

这将创建一个新的文件夹 esm_residues，其中包含 sampled_pdb 中每个 pdb 文件的 10 种潜在残基。

ProteinMPNN

要设置 ProteinMPNN，请参阅作者在其 GitHub 上提供的指南。

之后，我们遵循与 ESM-IF1 类似的流程（见上文），即 cd 到包含 sampled_pdb 文件夹的目录，并运行：

python ~/projects/foldingdiff/bin/pdb_to_residue_proteinmpnn.py sampled_pdb

这将创建一个名为 proteinmpnn_residues 的新目录，其中包含每条采样 PDB 结构对应的 8 条氨基酸链。

结构预测

生成氨基酸序列后，我们通过将其输入 OmegaFold 或 AlphaFold 来检查这些序列是否能够重现我们最初的采样结构。运行其中一个程序后，我们使用以下命令来评估自一致性 TM 分数：

python ~/projects/foldingdiff/bin/sctm.py -f alphafold_predictions_proteinmpnn

其中 alphafold_predictions_proteinmpnn 是一个包含反向折叠氨基酸序列对应折叠结构的文件夹。这将生成一个包含所有 scTM 分数的 JSON 文件，以及多个包含 scTM 分布图表和相关性的 PDF 文件。

OmegaFold

我们主要使用 OmegaFold 来折叠由 ESM-IF1 或 ProteinMPNN 生成的氨基酸序列。这是因为 OmegaFold 的运行时间相对 AlphaFold2 更短，而且 OmegaFold 原生设计为无需 MSA 信息即可运行——这使其更适合我们的蛋白质设计任务。

在创建并激活一个独立的 conda 环境，并按照作者的说明安装 OmegaFold 后，我们使用以下脚本将输入的氨基酸 FASTA 文件分配到多个 GPU 上进行推理，随后计算自一致性 TM (scTM) 分数。

# 将每个 FASTA 文件分散到 GPU 0 和 1 上进行折叠，输出到 omegafold_predictions 文件夹
python ~/projects/foldingdiff/bin/omegafold_across_gpus.py esm_residues/*.fasta -g 0 1

AlphaFold2

我们通过 localcolabfold 安装方式运行 AlphaFold2（详见 GitHub）。由于 AlphaFold 的运行时需求，我们提供了将 FASTA 文件集拆分为多个子目录的脚本，以便分别进行折叠；有关示例，请参阅 scripts/slurm/alphafold.sbatch 中的 SLURM 脚本。

测试

测试通过混合使用 doctest 和 unittest 实现。要运行 unittests，请执行以下命令：

python -m unittest -v

您可能会看到类似以下的警告；这是预期现象。

WARNING:root:protdiff-main/data/cath/dompdb/5a2qw00 中 omega 的值非法 -- 跳过

FoldingDiff 快速上手指南

FoldingDiff 是一个基于扩散模型（Diffusion Model）生成新型蛋白质骨架结构的开源工具。本指南将帮助开发者快速完成环境搭建并生成蛋白质结构。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux 或 macOS (Windows 需使用 WSL)
• 包管理器: 推荐安装 Conda 或 Mamba (Mamba 解析依赖速度更快)
• 硬件建议:
  • 训练/大量采样: 需要 NVIDIA GPU (支持 CUDA)
  • 简单采样: 可使用 CPU，但速度较慢（采样 512 个结构在 CPU 上约需 2 小时，而在 RTX 2080Ti 上仅需 7 分钟）
• 前置知识: 熟悉 Python 基础及 PyTorch 框架

安装步骤

1. 克隆代码库

git clone https://github.com/wukevin/foldingdiff.git
cd foldingdiff

2. 创建并激活虚拟环境

项目依赖定义在 environment.yml 中。推荐使用 mamba 加速环境创建过程：

# 使用 mamba (推荐)
mamba env create -f environment.yml
# 或者使用 conda
conda env create -f environment.yml

# 激活环境
conda activate foldingdiff

3. 安装项目包

确保在当前目录下执行以下命令以安装可编辑版本：

pip install -e ./

注意: 如果您仅需要进行**采样（生成结构）**而不需要重新训练模型，则无需下载额外的数据集。若需训练自己的模型，请参考原文档下载 CATH 数据集。

基本使用

方式一：命令行采样（推荐）

使用预训练模型快速生成蛋白质骨架。以下命令将生成长度在 50 到 128 之间的 10 个蛋白质骨架，批次大小为 512。

# 使用 GPU 运行 (默认 cuda:0)
python bin/sample.py -l 50 128 -n 10 -b 512 --device cuda:0

# 如果无 GPU，可使用 CPU 运行 (速度较慢)
# python bin/sample.py -l 50 128 -n 10 -b 512 --device cpu

输出说明： 运行完成后，当前目录将生成以下文件夹：

• sampled_pdb/: 包含生成的 .pdb 文件（笛卡尔坐标格式），可直接用于分子可视化软件（如 PyMOL）。
• sampled_angles/: 包含采样角度的 .csv.gz 文件。
• plots/: 包含训练数据与生成数据的角度分布对比图。
• model_snapshot/: 本次运行所使用的模型副本。

方式二：Python 代码调用

您也可以在 Python 脚本中直接加载预训练模型进行推理：

from huggingface_hub import snapshot_download
from torch.utils.data.dataloader import DataLoader
from foldingdiff import modelling
from foldingdiff import datasets as dsets

# 1. 加载预训练模型 (文件会自动缓存)
m = modelling.BertForDiffusion.from_dir(snapshot_download("wukevin/foldingdiff_cath"))

# 2. 准备数据集 (此处仅为演示加载流程，实际采样通常直接使用 sample.py)
clean_dset = dsets.CathCanonicalAnglesOnlyDataset(pad=128, trim_strategy='randomcrop')
noised_dset = dsets.NoisedAnglesDataset(clean_dset, timesteps=1000, beta_schedule='cosine')
dl = DataLoader(noised_dset, batch_size=32, shuffle=False)
x = iter(dl).next()

# 3. 执行前向传播
predicted_noise = m(x['corrupted'], x['t'], x['attn_mask'])

进阶：可视化折叠过程

若要生成展示扩散模型“折叠”过程的 GIF 动画，需在采样时添加 --fullhistory 参数，并使用 PyMOL 处理生成的历史帧：

1. 采样并记录历史:

   python bin/sample.py -l 50 128 -n 1 -b 1 --fullhistory --device cuda:0
   

2. 生成 GIF (需安装 PyMOL):

   python foldingdiff/pymol_vis.py pdb2gif -i sampled_pdb/sample_history/generated_0/*.pdb -o generated_0.gif