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链接中找到。
训练模型
要在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个模型和训练损失最佳的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
# 计算完这些之后,结果会保存在foldingdiff源目录下的.pkl文件中,以便将来更快地加载。
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
将默认使用第一个设备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笔记本。我们发现以下一系列命令在我们的机器上有效:
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