alphafold3-pytorch

Alphafold 3 - Pytorch

Pytorch 实现的 <a href="https://www.nature.com/articles/s41586-024-07487-w">Alphafold 3</a>

您可以在这里与其他研究人员讨论这项工作

<a href="https://www.youtube.com/watch?v=qjFgthkKxcA">论文评述</a>由<a href="https://x.com/sokrypton">Sergey</a>提供

<a href="https://elanapearl.github.io/blog/2024/the-illustrated-alphafold/">图解指南</a>由<a href="https://elanapearl.github.io/">Elana P. Simon</a>提供

<a href="https://www.youtube.com/watch?v=AE35XCN5NuU">Max Jaderberg 的演讲</a>

<a href="https://github.com/amorehead">Alex</a>在<a href="https://github.com/amorehead/alphafold3-pytorch-lightning-hydra">这个仓库</a>中维护了一个完整支持 Lightning + Hydra 的分支

可以在<a href="https://colab.research.google.com/drive/1S9TD8WmS1Gjgwjo9xyEYTbdxgMVVZcQe?usp=sharing">这里</a>查看并交互使用仓库中生命分子的可视化

致谢

<a href="https://github.com/joseph-c-kim">Joseph</a> 贡献了相对位置编码和平滑 LDDT 损失函数！
<a href="https://github.com/engelberger">Felipe</a> 贡献了加权刚性对齐、在框架中表达坐标、计算对齐误差和中心随机增强模块！
<a href="https://github.com/amorehead">Alex</a> 修复了转录算法中的各种问题
<a href="https://github.com/gitabtion">Heng</a> 指出了与论文不一致的地方并提交了解决方案
<a href="https://github.com/gitabtion">Heng</a> 发现了距离图损失中分子原子索引的问题
<a href="https://github.com/luwei0917">Wei Lu</a> 发现了一些错误的超参数
<a href="https://github.com/amorehead">Alex</a> 提供了 PDB 数据集准备脚本！
<a href="https://github.com/milot-mirdita">Milot</a> 优化了 PDB 数据集聚类脚本！
<a href="https://github.com/amorehead">Alex</a> 基本上编写了整个从解析 PDB 到训练所需分子和原子输入的庞大流程
<a href="https://github.com/vandrw">Andrei</a> 致力于加权 PDB 数据集采样！
<a href="https://github.com/tanjimin">Jimin</a> 提交了一个小修复，解决了传入 WeightedRigidAlign 的坐标问题
<a href="https://github.com/xluo233">@xluo233</a> 贡献了置信度度量、碰撞惩罚排序和样本排序逻辑！
<a href="https://github.com/sj900">sj900</a> 在 PDBDataset 中集成并测试了 WeightedPDBSampler！
<a href="https://github.com/xluo233">@xluo233</a> 再次贡献了计算模型选择分数以及未解析 rasa 的逻辑！
<a href="https://github.com/wufandi">Fandi</a> 发现了阐明的原子扩散模块与补充材料中的几处不一致之处
<a href="https://github.com/ptosco">Paolo</a> 提出了 PDB 中性稳定分子 假设！
<a href="https://github.com/dhuvik">Dhuvi</a> 修复了与 Alphafold3Inputs 中金属离子分子 ID 分配相关的错误！
感谢 <a href="https://github.com/patrick-kidger">Patrick</a> 的 <a href="https://docs.kidger.site/jaxtyping/">jaxtyping</a>，<a href="https://github.com/fferflo">Florian</a> 的 <a href="https://github.com/fferflo/einx">einx</a>，当然还有 <a href="https://github.com/arogozhnikov">Alex</a> 的 <a href="https://einops.rocks/">einops</a>

安装

$ pip install alphafold3-pytorch

使用方法

import torch
from alphafold3_pytorch import Alphafold3

alphafold3 = Alphafold3(
    dim_atom_inputs = 77,
    dim_template_feats = 44
)

# 模拟输入

seq_len = 16
molecule_atom_lens = torch.randint(1, 3, (2, seq_len))
atom_seq_len = molecule_atom_lens.sum(dim = -1).amax()

atom_inputs = torch.randn(2, atom_seq_len, 77)
atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)

additional_molecule_feats = torch.randint(0, 2, (2, seq_len, 5))
additional_token_feats = torch.randn(2, seq_len, 2)
is_molecule_types = torch.randint(0, 2, (2, seq_len, 5)).bool()
is_molecule_mod = torch.randint(0, 2, (2, seq_len, 4)).bool()
molecule_ids = torch.randint(0, 32, (2, seq_len))

template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
template_mask = torch.ones((2, 2)).bool()

msa = torch.randn(2, 7, seq_len, 64)
msa_mask = torch.ones((2, 7)).bool()

# 训练时需要，推理时可省略

atom_pos = torch.randn(2, atom_seq_len, 3)
molecule_atom_indices = molecule_atom_lens - 1 # 最后一个原子，作为示例

distance_labels = torch.randint(0, 37, (2, seq_len, seq_len))
pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
pde_labels = torch.randint(0, 64, (2, seq_len, seq_len))
plddt_labels = torch.randint(0, 50, (2, seq_len))
resolved_labels = torch.randint(0, 2, (2, seq_len))

# 训练

loss = alphafold3(
    num_recycling_steps = 2,
    atom_inputs = atom_inputs,
    atompair_inputs = atompair_inputs,
    molecule_ids = molecule_ids,
    molecule_atom_lens = molecule_atom_lens,
    additional_molecule_feats = additional_molecule_feats,
    additional_token_feats = additional_token_feats,
    is_molecule_types = is_molecule_types,
    is_molecule_mod = is_molecule_mod,
    msa = msa,
    msa_mask = msa_mask,
    templates = template_feats,
    template_mask = template_mask,
    atom_pos = atom_pos,
    molecule_atom_indices = molecule_atom_indices,
    distance_labels = distance_labels,
    pae_labels = pae_labels,
    pde_labels = pde_labels,
    plddt_labels = plddt_labels,
    resolved_labels = resolved_labels
)

loss.backward()

# 经过大量训练后...

sampled_atom_pos = alphafold3(
    num_recycling_steps = 4,
    num_sample_steps = 16,
    atom_inputs = atom_inputs,
    atompair_inputs = atompair_inputs,
    molecule_ids = molecule_ids,
    molecule_atom_lens = molecule_atom_lens,
    additional_molecule_feats = additional_molecule_feats,
    additional_token_feats = additional_token_feats,
    is_molecule_types = is_molecule_types,
    is_molecule_mod = is_molecule_mod,
    msa = msa,
    msa_mask = msa_mask,
    templates = template_feats,
    template_mask = template_mask
)

sampled_atom_pos.shape # (2, <atom_seqlen>, 3)

一个分子级输入处理的示例

import torch

from alphafold3_pytorch import (
    Alphafold3,
    Alphafold3Input,
    alphafold3_inputs_to_batched_atom_input
)

contrived_protein = 'AG'

mock_atompos = [
    torch.randn(5, 3),   # 丙氨酸有 5 个非氢原子
    torch.randn(4, 3)    # 甘氨酸有 4 个非氢原子
]

train_alphafold3_input = Alphafold3Input(
    proteins = [contrived_protein],
    atom_pos = mock_atompos
)

eval_alphafold3_input = Alphafold3Input(
    proteins = [contrived_protein]
)

batched_atom_input = alphafold3_inputs_to_batched_atom_input(train_alphafold3_input, atoms_per_window = 27)

# 训练

alphafold3 = Alphafold3(
    dim_atom_inputs = 3,
    dim_atompair_inputs = 5,
    atoms_per_window = 27,
    dim_template_feats = 44,
    num_dist_bins = 38,
    num_molecule_mods = 0,
    confidence_head_kwargs = dict(
        pairformer_depth = 1
    ),
    template_embedder_kwargs = dict(
        pairformer_stack_depth = 1
    ),
    msa_module_kwargs = dict(
        depth = 1
    ),
    pairformer_stack = dict(
        depth = 2
    ),
    diffusion_module_kwargs = dict(
        atom_encoder_depth = 1,
        token_transformer_depth = 1,
        atom_decoder_depth = 1,
    )
)

loss = alphafold3(**batched_atom_input.model_forward_dict())
loss.backward()

# 采样

batched_eval_atom_input = alphafold3_inputs_to_batched_atom_input(eval_alphafold3_input, atoms_per_window = 27)
alphafold3.eval()
sampled_atom_pos = alphafold3(**batched_eval_atom_input.model_forward_dict())

assert sampled_atom_pos.shape == (1, (5 + 4), 3)

数据准备

PDB数据集整理

要获取AlphaFold 3 PDB数据集，首先下载蛋白质数据库（PDB）中的所有第一组装（和不对称单元）复合物，然后使用下面引用的脚本对它们进行预处理。PDB可以从RCSB下载：https://www.wwpdb.org/ftp/pdb-ftp-sites#rcsbpdb。以下两个Python脚本（即`filter_pdb_{train,val,test}_mmcifs.py`和`cluster_pdb_{train,val,test}_mmcifs.py`）假设你已经以**mmCIF文件格式**下载了PDB，并将其第一组装和不对称单元mmCIF文件分别放置在`data/pdb_data/unfiltered_assembly_mmcifs/`和`data/pdb_data/unfiltered_asym_mmcifs/`。

为了可重现性，我们建议使用AWS快照（例如20240101）下载PDB。要做到这一点，请参考AWS的文档在本地设置AWS CLI。或者，在RCSB网站上，向下滚动到"下载协议"，并根据你的位置按照下载说明进行操作。

例如，可以使用以下命令将PDB下载为两个mmCIF文件集合：

# 对于`assembly1`复合物，使用PDB的`20240101` AWS快照：
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs
# 或者作为备选方案，使用rsync：
rsync -rlpt -v -z --delete --port=33444 \
rsync.rcsb.org::ftp_data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs/

# 对于不对称单元复合物，同样使用PDB的`20240101` AWS快照：
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs
# 或者作为备选方案，使用rsync：
rsync -rlpt -v -z --delete --port=33444 \
rsync.rcsb.org::ftp_data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs/

警告：下载PDB可能会占用高达700GB的空间。

注意：PDB在这里托管所有可用的AWS快照：https://pdbsnapshots.s3.us-west-2.amazonaws.com/index.html。

下载后，你应该有两个格式如下的目录： https://files.rcsb.org/pub/pdb/data/assemblies/mmCIF/divided/ 和 https://files.rcsb.org/pub/pdb/data/structures/divided/mmCIF/

00/
01/
02/
..
zz/

对于这些目录，解压所有文件：

find ./data/pdb_data/unfiltered_assembly_mmcifs/ -type f -name "*.gz" -exec gzip -d {} \;
find ./data/pdb_data/unfiltered_asym_mmcifs/ -type f -name "*.gz" -exec gzip -d {} \;

接下来在项目根目录运行以下命令

wget -P ./data/ccd_data/ https://files.wwpdb.org/pub/pdb/data/monomers/components.cif.gz
wget -P ./data/ccd_data/ https://files.wwpdb.org/pub/pdb/data/component-models/complete/chem_comp_model.cif.gz

以下载PDB化学成分字典（CCD）的最新版本及其结构模型。使用以下命令解压每个文件：

find data/ccd_data/ -type f -name "*.gz" -exec gzip -d {} \;

PDB数据集过滤

然后运行以下命令，将pdb_assembly_dir、pdb_asym_dir、ccd_dir和mmcif_output_dir替换为你本地第一组装PDB、不对称单元PDB、CCD的副本位置以及你想要的数据集输出目录（即./data/pdb_data/unfiltered_assembly_mmcifs/、./data/pdb_data/unfiltered_asym_mmcifs/、./data/ccd_data/和./data/pdb_data/{train,val,test}_mmcifs/）。

python scripts/filter_pdb_train_mmcifs.py --mmcif_assembly_dir <pdb_assembly_dir> --mmcif_asym_dir <pdb_asym_dir> --ccd_dir <ccd_dir> --output_dir <mmcif_output_dir>
python scripts/filter_pdb_val_mmcifs.py --mmcif_assembly_dir <pdb_assembly_dir> --mmcif_asym_dir <pdb_asym_dir> --output_dir <mmcif_output_dir>
python scripts/filter_pdb_test_mmcifs.py --mmcif_assembly_dir <pdb_assembly_dir> --mmcif_asym_dir <pdb_asym_dir> --output_dir <mmcif_output_dir>

查看脚本以了解更多选项。每个成功通过所有处理步骤的第一组装mmCIF将被写入mmcif_output_dir中的一个子目录，该子目录根据mmCIF的第二个和第三个PDB ID字符命名（例如5c）。

PDB数据集聚类

接下来，运行以下命令，将mmcif_dir和{train,val,test}_clustering_output_dir分别替换为使用上述数据集过滤脚本创建的本地输出目录和你想要的聚类输出目录（即./data/pdb_data/{train,val,test}_mmcifs/和./data/pdb_data/data_caches/{train,val,test}_clusterings/）：

python scripts/cluster_pdb_train_mmcifs.py --mmcif_dir <mmcif_dir> --output_dir <train_clustering_output_dir> --clustering_filtered_pdb_dataset
python scripts/cluster_pdb_val_mmcifs.py --mmcif_dir <mmcif_dir> --reference_clustering_dir <train_clustering_output_dir> --output_dir <val_clustering_output_dir> --clustering_filtered_pdb_dataset
python scripts/cluster_pdb_test_mmcifs.py --mmcif_dir <mmcif_dir> --reference_1_clustering_dir <train_clustering_output_dir> --reference_2_clustering_dir <val_clustering_output_dir> --output_dir <test_clustering_output_dir> --clustering_filtered_pdb_dataset

注意：在使用上述脚本整理的过滤后PDB数据集进行聚类时，建议使用--clustering_filtered_pdb_dataset标志，因为在这种情况下，此标志将实现更快的运行时间（因为过滤后每个链的残基ID都是从1开始的）。但是，在对其他（即非PDB）mmCIF文件数据集进行聚类时，必须不提供此标志。否则，界面聚类可能会执行错误，因为这些数据集的mmCIF文件可能不会为每个链使用严格的从1开始的残基索引。

注意：也可以通过共享的OneDrive文件夹下载PDB 20240101 AWS快照的预处理（即过滤后）mmCIF（train/val/test）文件（约25GB，包含148k个复合物）和链/界面聚类（train/val/test）文件（约3GB）。这些tar压缩包应在data/pdb_data/目录下解压。

贡献

在项目根目录运行

$ sh ./contribute.sh

然后，将你的模块添加到alphafold3_pytorch/alphafold3.py，将你的测试添加到tests/test_af3.py，并提交一个拉取请求。你可以通过以下命令在本地运行测试：

$ pytest tests/

Docker镜像

包含的Dockerfile包含运行该包所需的依赖项，以及使用PyTorch和GPU进行训练/推理所需的依赖项。

默认的基础镜像是pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime，并从GitHub的main分支安装此包的最新版本。

## 构建Docker容器
docker build -t af3 .

或者，使用构建参数重新构建镜像，以使用不同的软件版本：

PYTORCH_TAG：更改基础镜像，从而使用不同版本的PyTorch、CUDA和/或cuDNN进行构建。
GIT_TAG：更改要克隆和安装的此仓库的标签。

例如：

## 使用构建参数更改版本
docker build --build-arg "PYTORCH_TAG=2.2.1-cuda12.1-cudnn8-devel" --build-arg "GIT_TAG=0.1.15" -t af3 .

然后，使用以下命令运行带有GPU的容器并挂载本地卷（用于训练）：

## 运行容器
docker run -v .:/data --gpus all -it af3

引用

@article{Abramson2024-fj,
  标题    = "使用AlphaFold 3准确预测生物分子相互作用的结构",
  作者    = "Abramson, Josh 和 Adler, Jonas 和 Dunger, Jack 和 Evans,
              Richard 和 Green, Tim 和 Pritzel, Alexander 和 Ronneberger,
              Olaf 和 Willmore, Lindsay 和 Ballard, Andrew J 和 Bambrick,
              Joshua 和 Bodenstein, Sebastian W 和 Evans, David A 和 Hung,
              Chia-Chun 和 O'Neill, Michael 和 Reiman, David 和
              Tunyasuvunakool, Kathryn 和 Wu, Zachary 和 Žemgulytė,
              Akvilė 和 Arvaniti, Eirini 和 Beattie, Charles 和
              Bertolli, Ottavia 和 Bridgland, Alex 和 Cherepanov, Alexey 和
              Congreve, Miles 和 Cowen-Rivers, Alexander I 和 Cowie, Andrew
              和 Figurnov, Michael 和 Fuchs, Fabian B 和 Gladman, Hannah 和
              Jain, Rishub 和 Khan, Yousuf A 和 Low, Caroline M R 和 Perlin,
              Kuba 和 Potapenko, Anna 和 Savy, Pascal 和 Singh, Sukhdeep 和
              Stecula, Adrian 和 Thillaisundaram, Ashok 和 Tong, Catherine
              和 Yakneen, Sergei 和 Zhong, Ellen D 和 Zielinski, Michal 和
              Žídek, Augustin 和 Bapst, Victor 和 Kohli, Pushmeet
              和 Jaderberg, Max 和 Hassabis, Demis 和 Jumper, John M",
  期刊    = "《自然》",
  月份    = "五月",
  年份    =  2024
}

@inproceedings{Darcet2023VisionTN,
    标题   = {视觉transformer需要寄存器},
    作者   = {Timothée Darcet 和 Maxime Oquab 和 Julien Mairal 和 Piotr Bojanowski},
    年份   = {2023},
    网址   = {https://api.semanticscholar.org/CorpusID:263134283}
}

@article{Arora2024SimpleLA,
    标题   = {简单线性注意力语言模型在回忆和吞吐量之间取得平衡},
    作者   = {Simran Arora 和 Sabri Eyuboglu 和 Michael Zhang 和 Aman Timalsina 和 Silas Alberti 和 Dylan Zinsley 和 James Zou 和 Atri Rudra 和 Christopher Ré},
    期刊   = {ArXiv},
    年份   = {2024},
    卷号   = {abs/2402.18668},
    网址   = {https://api.semanticscholar.org/CorpusID:268063190}
}

@article{Puny2021FrameAF,
    标题   = {帧平均用于不变和等变网络设计},
    作者   = {Omri Puny 和 Matan Atzmon 和 Heli Ben-Hamu 和 Edward James Smith 和 Ishan Misra 和 Aditya Grover 和 Yaron Lipman},
    期刊   = {ArXiv},
    年份   = {2021},
    卷号   = {abs/2110.03336},
    网址   = {https://api.semanticscholar.org/CorpusID:238419638}
}

@article{Duval2023FAENetFA,
    标题   = {FAENet：用于材料建模的帧平均等变图神经网络},
    作者   = {Alexandre Duval 和 Victor Schmidt 和 Alex Hernandez Garcia 和 Santiago Miret 和 Fragkiskos D. Malliaros 和 Yoshua Bengio 和 David Rolnick},
    期刊   = {ArXiv},
    年份   = {2023},
    卷号   = {abs/2305.05577},
    网址   = {https://api.semanticscholar.org/CorpusID:258564608}
}

@article{Wang2022DeepNetST,
    标题   = {DeepNet：将Transformer扩展到1,000层},
    作者   = {Hongyu Wang 和 Shuming Ma 和 Li Dong 和 Shaohan Huang 和 Dongdong Zhang 和 Furu Wei},
    期刊   = {ArXiv},
    年份   = {2022},
    卷号   = {abs/2203.00555},
    网址   = {https://api.semanticscholar.org/CorpusID:247187905}
}

@inproceedings{Ainslie2023CoLT5FL,
    标题   = {CoLT5：使用条件计算的更快长程Transformer},
    作者   = {Joshua Ainslie 和 Tao Lei 和 Michiel de Jong 和 Santiago Ontañón 和 Siddhartha Brahma 和 Yury Zemlyanskiy 和 David Uthus 和 Mandy Guo 和 James Lee-Thorp 和 Yi Tay 和 Yun-Hsuan Sung 和 Sumit Sanghai},
    年份   = {2023}
}