альфафолд3
0.6.6
Реализация Alphafold 3 в Pytorch
Вы можете пообщаться с другими исследователями об этой работе здесь.
Рецензия на статью Сергея
Иллюстрированное руководство Эланы П. Саймон.
Выступление Макса Ядерберга
Форк с полной поддержкой Lightning + Hydra поддерживается Алексом в этом репозитории.
Визуализацию молекул жизни, используемых в хранилище, можно увидеть и поработать здесь.
Джозефу за вклад в относительное позиционное кодирование и плавную потерю LDDT!
Фелипе за участие в модулях «Взвешенное жесткое выравнивание», «Экспресс-координаты в кадре», «Вычисление ошибки выравнивания» и «Случайное увеличение центра»!
Алексу за исправление различных ошибок в записанных алгоритмах.
Хэну за указание на несоответствия в статье и запрос решений
Хэну за обнаружение проблемы с индексами молекулярных атомов, приводящей к потере дистограммы.
Вэй Лу за обнаружение нескольких ошибочных гиперпараметров.
Алексу за скрипт подготовки набора данных PDB!
Майлоту за оптимизацию скрипта кластеризации набора данных PDB!
Алексу за написание всего гигантского потока от анализа PDB до молекул и атомарных входных данных для обучения.
Андрею за работу над взвешенной выборкой набора данных PDB!
Чимину за небольшое исправление проблемы с передачей координат в WeightedRigidAlign
@xluo233 за предоставленные меры достоверности, ранжирование штрафов за коллизии и пример логики ранжирования!
sj900 за интеграцию и тестирование WeightedPDBSampler в PDBDataset , а также за добавление начальной поддержки MSA и анализа шаблонов!
@xluo233 еще раз за предоставленную логику для расчета оценки выбора модели, а также неразрешенной расы!
Фанди за обнаружение нескольких несоответствий в объясненном модуле диффузии атомов с дополнительным
Паоло за предложение гипотезы PDB neutral stable molecule !
Дхуви за исправление ошибки, связанной с присвоением идентификатора молекулы иона металла для Alphafold3Inputs !
Дхуви за то, что взял на себя логику перевода Alphafold3Input в BioMolecule для сохранения в mmCIF!
Тому (с канала Discord) за обнаружение несоответствия между вычислениями дистограммы и единичного вектора шаблона в этой кодовой базе и вычислениями OpenFold (и Андрею за помощь в решении проблемы с дистограммой)!
Кайхуи за обнаружение ошибки в обработке нестандартных атомов в остатках полимера!
Андрею за работу над градиентным интерфейсом!
Патрик за набор текста, Флориан за einx и, конечно же, Алекс за einops.
Сумиту и организации Pytorch за предоставленную мне возможность открыть исходный код этой работы.
$ pip install alphafold3-pytorch import torch
from alphafold3_pytorch import Alphafold3
from alphafold3_pytorch . utils . model_utils import exclusive_cumsum
alphafold3 = Alphafold3 (
dim_atom_inputs = 77 ,
dim_template_feats = 108
)
# mock inputs
seq_len = 16
molecule_atom_indices = torch . randint ( 0 , 2 , ( 2 , seq_len )). long ()
molecule_atom_lens = torch . full (( 2 , seq_len ), 2 ). long ()
atom_seq_len = molecule_atom_lens . sum ( dim = - 1 ). amax ()
atom_offsets = exclusive_cumsum ( molecule_atom_lens )
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 , 33 )
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 , 108 )
template_mask = torch . ones (( 2 , 2 )). bool ()
msa = torch . randn ( 2 , 7 , seq_len , 32 )
msa_mask = torch . ones (( 2 , 7 )). bool ()
additional_msa_feats = torch . randn ( 2 , 7 , seq_len , 2 )
# required for training, but omitted on inference
atom_pos = torch . randn ( 2 , atom_seq_len , 3 )
distogram_atom_indices = molecule_atom_lens - 1
distance_labels = torch . randint ( 0 , 37 , ( 2 , seq_len , seq_len ))
resolved_labels = torch . randint ( 0 , 2 , ( 2 , atom_seq_len ))
# offset indices correctly
distogram_atom_indices += atom_offsets
molecule_atom_indices += atom_offsets
# train
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_msa_feats = additional_msa_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 ,
distogram_atom_indices = distogram_atom_indices ,
molecule_atom_indices = molecule_atom_indices ,
distance_labels = distance_labels ,
resolved_labels = resolved_labels
)
loss . backward ()
# after much training ...
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_msa_feats = additional_msa_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, , 3) Пример обработки входных данных на уровне молекул
import torch
from alphafold3_pytorch import Alphafold3 , Alphafold3Input
contrived_protein = 'AG'
mock_atompos = [
torch . randn ( 5 , 3 ), # alanine has 5 non-hydrogen atoms
torch . randn ( 4 , 3 ) # glycine has 4 non-hydrogen atoms
]
train_alphafold3_input = Alphafold3Input (
proteins = [ contrived_protein ],
atom_pos = mock_atompos
)
eval_alphafold3_input = Alphafold3Input (
proteins = [ contrived_protein ]
)
# training
alphafold3 = Alphafold3 (
dim_atom_inputs = 3 ,
dim_atompair_inputs = 5 ,
atoms_per_window = 27 ,
dim_template_feats = 108 ,
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 . forward_with_alphafold3_inputs ([ train_alphafold3_input ])
loss . backward ()
# sampling
alphafold3 . eval ()
sampled_atom_pos = alphafold3 . forward_with_alphafold3_inputs ( eval_alphafold3_input )
assert sampled_atom_pos . shape == ( 1 , ( 5 + 4 ), 3 ) Чтобы получить набор данных 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 ) предполагают, что вы загрузили PDB в формате файла mmCIF , разместив его первую сборку и файлы mmCIF с асимметричными модулями по адресу data/pdb_data/unfiltered_assembly_mmcifs/ и data/pdb_data/unfiltered_asym_mmcifs/ соответственно.
Для воспроизводимости мы рекомендуем загружать PDB с помощью снимков AWS (например, 20240101 ). Для этого обратитесь к документации AWS, чтобы настроить AWS CLI локально. Альтернативно, на веб-сайте RCSB перейдите к разделу «Протоколы загрузки» и следуйте инструкциям по загрузке в зависимости от вашего местоположения.
Например, можно использовать следующие команды для загрузки PDB в виде двух коллекций файлов mmCIF:
# For `assembly1` complexes, use the PDB's `20240101` AWS snapshot:
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs
# Or as a fallback, use rsync:
rsync -rlpt -v -z --delete --port=33444
rsync.rcsb.org::ftp_data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs/
# For asymmetric unit complexes, also use the PDB's `20240101` AWS snapshot:
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs
# Or as a fallback, use rsync:
rsync -rlpt -v -z --delete --port=33444
rsync.rcsb.org::ftp_data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs/ВНИМАНИЕ: Загрузка PDB может занять до 700 ГБ места.
ПРИМЕЧАНИЕ. Все доступные снимки AWS PDB размещаются здесь: 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. /структуры/разделенные/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из корневого каталога проекта, чтобы загрузить последнюю версию Словаря химических компонентов (CCD) PDB и его структурных моделей. Извлеките каждый из этих файлов с помощью следующей команды:
find data/ccd_data/ -type f -name " *.gz " -exec gzip -d {} ; Затем запустите следующую команду, заменив 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 в подкаталоге, названном в соответствии со вторым и третьим символами идентификатора PDB mmCIF (например, 5c ).
Затем запустите следующую команду, заменив 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 Примечание . Флаг --clustering_filtered_pdb_dataset рекомендуется использовать при кластеризации отфильтрованного набора данных PDB, как рекомендовано, с использованием приведенных выше сценариев, поскольку этот флаг обеспечит более быстрое время выполнения в этом контексте (поскольку фильтрация оставляет идентификаторы остатков каждой цепочки основанными на 1). Однако этот флаг не следует указывать при кластеризации других (т. е. не PDB) наборов данных файлов mmCIF. В противном случае кластеризация интерфейса может выполняться неправильно, поскольку файлы mmCIF этих наборов данных могут не использовать строгую индексацию остатков на основе 1 для каждой цепочки.
Примечание . Вместо этого можно загрузить предварительно обработанные (т. е. отфильтрованные) файлы mmCIF ( train / val / test ) (~25 ГБ, включая 148 тыс. комплексов) и файлы цепочки/интерфейсной кластеризации ( train / val / test ) (~3 ГБ) для PDB 20240101 Снимок AWS через общую папку OneDrive. Каждый из этих архивов tar.gz следует распаковать в каталоге data/pdb_data/ например, с помощью tar -xzf data_caches.tar.gz -C data/pdb_data/ . Также можно загрузить и подготовить данные дистилляции PDB, используя в качестве ссылки сценарий scripts/distillation_data_download.sh . После загрузки можно запустить scripts/reduce_uniprot_predictions_to_pdb.py , чтобы отфильтровать этот набор данных только для примеров, связанных хотя бы с одной записью PDB. Более того, для удобства сопоставление идентификаторов доступа UniProt с идентификаторами PDB для обучения на данных дистилляции PDB уже загружено и извлечено как data/afdb_data/data_caches/uniprot_to_pdb_id_mapping.dat .
В корне проекта запустите
$ sh ./contribute.sh Затем добавьте свой модуль в alphafold3_pytorch/alphafold3.py , добавьте свои тесты tests/test_af3.py и отправьте запрос на включение. Вы можете запускать тесты локально с помощью
$ pytest tests/ Включенный Dockerfile содержит необходимые зависимости для запуска пакета и обучения/вывода с использованием PyTorch с графическими процессорами.
Базовый образ по умолчанию — pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime , он устанавливает последнюю версию этого пакета из main ветки GitHub.
# # Build Docker Container
docker build -t af3 .Альтернативно, используйте аргументы сборки, чтобы перестроить образ с разными версиями программного обеспечения:
PYTORCH_TAG : изменяет базовый образ и, таким образом, строит его с использованием разных версий PyTorch, CUDA и/или cuDNN.GIT_TAG : изменяет тег этого репозитория для клонирования и установки пакета.Например:
# # Use build argument to change versions
docker build --build-arg " PYTORCH_TAG=2.2.1-cuda12.1-cudnn8-devel " --build-arg " GIT_TAG=0.1.15 " -t af3 .Затем запустите контейнер с графическими процессорами и смонтируйте локальный том (для обучения) с помощью следующей команды:
# # Run Container
docker run -v .:/data --gpus all -it af3 @article { Abramson2024-fj ,
title = " Accurate structure prediction of biomolecular interactions with
{AlphaFold} 3 " ,
author = " Abramson, Josh and Adler, Jonas and Dunger, Jack and Evans,
Richard and Green, Tim and Pritzel, Alexander and Ronneberger,
Olaf and Willmore, Lindsay and Ballard, Andrew J and Bambrick,
Joshua and Bodenstein, Sebastian W and Evans, David A and Hung,
Chia-Chun and O'Neill, Michael and Reiman, David and
Tunyasuvunakool, Kathryn and Wu, Zachary and {v Z}emgulyt{.e},
Akvil{.e} and Arvaniti, Eirini and Beattie, Charles and
Bertolli, Ottavia and Bridgland, Alex and Cherepanov, Alexey and
Congreve, Miles and Cowen-Rivers, Alexander I and Cowie, Andrew
and Figurnov, Michael and Fuchs, Fabian B and Gladman, Hannah and
Jain, Rishub and Khan, Yousuf A and Low, Caroline M R and Perlin,
Kuba and Potapenko, Anna and Savy, Pascal and Singh, Sukhdeep and
Stecula, Adrian and Thillaisundaram, Ashok and Tong, Catherine
and Yakneen, Sergei and Zhong, Ellen D and Zielinski, Michal and
{v Z}{'i}dek, Augustin and Bapst, Victor and Kohli, Pushmeet
and Jaderberg, Max and Hassabis, Demis and Jumper, John M " ,
journal = " Nature " ,
month = " May " ,
year = 2024
} @inproceedings { Darcet2023VisionTN ,
title = { Vision Transformers Need Registers } ,
author = { Timoth'ee Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski } ,
year = { 2023 } ,
url = { https://api.semanticscholar.org/CorpusID:263134283 }
} @article { Arora2024SimpleLA ,
title = { Simple linear attention language models balance the recall-throughput tradeoff } ,
author = { Simran Arora and Sabri Eyuboglu and Michael Zhang and Aman Timalsina and Silas Alberti and Dylan Zinsley and James Zou and Atri Rudra and Christopher R'e } ,
journal = { ArXiv } ,
year = { 2024 } ,
volume = { abs/2402.18668 } ,
url = { https://api.semanticscholar.org/CorpusID:268063190 }
} @article { Puny2021FrameAF ,
title = { Frame Averaging for Invariant and Equivariant Network Design } ,
author = { Omri Puny and Matan Atzmon and Heli Ben-Hamu and Edward James Smith and Ishan Misra and Aditya Grover and Yaron Lipman } ,
journal = { ArXiv } ,
year = { 2021 } ,
volume = { abs/2110.03336 } ,
url = { https://api.semanticscholar.org/CorpusID:238419638 }
} @article { Duval2023FAENetFA ,
title = { FAENet: Frame Averaging Equivariant GNN for Materials Modeling } ,
author = { Alexandre Duval and Victor Schmidt and Alex Hernandez Garcia and Santiago Miret and Fragkiskos D. Malliaros and Yoshua Bengio and David Rolnick } ,
journal = { ArXiv } ,
year = { 2023 } ,
volume = { abs/2305.05577 } ,
url = { https://api.semanticscholar.org/CorpusID:258564608 }
} @article { Wang2022DeepNetST ,
title = { DeepNet: Scaling Transformers to 1, 000 Layers } ,
author = { Hongyu Wang and Shuming Ma and Li Dong and Shaohan Huang and Dongdong Zhang and Furu Wei } ,
journal = { ArXiv } ,
year = { 2022 } ,
volume = { abs/2203.00555 } ,
url = { https://api.semanticscholar.org/CorpusID:247187905 }
} @inproceedings { Ainslie2023CoLT5FL ,
title = { CoLT5: Faster Long-Range Transformers with Conditional Computation } ,
author = { Joshua Ainslie and Tao Lei and Michiel de Jong and Santiago Ontan'on and Siddhartha Brahma and Yury Zemlyanskiy and David Uthus and Mandy Guo and James Lee-Thorp and Yi Tay and Yun-Hsuan Sung and Sumit Sanghai } ,
year = { 2023 }
} @article { Ash2019OnTD ,
title = { On the Difficulty of Warm-Starting Neural Network Training } ,
author = { Jordan T. Ash and Ryan P. Adams } ,
journal = { ArXiv } ,
year = { 2019 } ,
volume = { abs/1910.08475 } ,
url = { https://api.semanticscholar.org/CorpusID:204788802 }
} @ARTICLE { Heinzinger2023.07.23.550085 ,
author = { Michael Heinzinger and Konstantin Weissenow and Joaquin Gomez Sanchez and Adrian Henkel and Martin Steinegger and Burkhard Rost } ,
title = { ProstT5: Bilingual Language Model for Protein Sequence and Structure } ,
year = { 2023 } ,
doi = { 10.1101/2023.07.23.550085 } ,
journal = { bioRxiv }
} @article { Lin2022.07.20.500902 ,
author = { Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Santos Costa, Allan dos and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Sal and Rives, Alexander } ,
title = { Language models of protein sequences at the scale of evolution enable accurate structure prediction } ,
elocation-id = { 2022.07.20.500902 } ,
year = { 2022 } ,
doi = { 10.1101/2022.07.20.500902 } ,
publisher = { Cold Spring Harbor Laboratory } ,
URL = { https://www.biorxiv.org/content/early/2022/07/21/2022.07.20.500902 } ,
eprint = { https://www.biorxiv.org/content/early/2022/07/21/2022.07.20.500902.full.pdf } ,
journal = { bioRxiv }
} @article { Li2024SwitchEA ,
title = { Switch EMA: A Free Lunch for Better Flatness and Sharpness } ,
author = { Siyuan Li and Zicheng Liu and Juanxi Tian and Ge Wang and Zedong Wang and Weiyang Jin and Di Wu and Cheng Tan and Tao Lin and Yang Liu and Baigui Sun and Stan Z. Li } ,
journal = { ArXiv } ,
year = { 2024 } ,
volume = { abs/2402.09240 } ,
url = { https://api.semanticscholar.org/CorpusID:267657558 }
} @article { Nguyen2023MitigatingOI ,
title = { Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals } ,
author = { Tam Nguyen and Tan M. Nguyen and Richard G. Baraniuk } ,
journal = { ArXiv } ,
year = { 2023 } ,
volume = { abs/2312.00751 } ,
url = { https://api.semanticscholar.org/CorpusID:264300597 }
} @inproceedings { Zhou2024ValueRL ,
title = { Value Residual Learning For Alleviating Attention Concentration In Transformers } ,
author = { Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan } ,
year = { 2024 } ,
url = { https://api.semanticscholar.org/CorpusID:273532030 }
}
