repository·main·Indexed 23 days ago
https://github.com/lucidrains/alphafold3-pytorchA PyTorch implementation of AlphaFold 3 for protein and molecule structure prediction. The library provides the Alphafold3 model class, high-level input abstractions via Alphafold3Input, and a comprehensive suite of architectural components including MSAModule, PairformerStack, and DiffusionTransformer. It includes utilities for PDB mmCIF dataset downloading, filtering, and clustering, as well as training tools like the Trainer class and SmoothLDDTLoss.
To contribute new modules or tests:
sh ./contribute.sh.alphafold3_pytorch/alphafold3.py.tests/test_af3.py.pytest tests/.$ sh ./contribute.sh
# Run tests locally
$ pytest tests/Cluster the filtered mmCIF files to prepare for training.
Important Flag: Use --clustering_filtered_pdb_dataset when clustering the curated PDB dataset. This enables faster runtimes because the filtering process ensures residue IDs are 1-based. Do not use this flag for non-PDB datasets, as it may cause incorrect interface clustering if the files do not use strict 1-based indexing.
Arguments:
--mmcif_dir: The directory containing filtered mmCIF files.--output_dir: The directory for clustering output.--reference_clustering_dir / --reference_1_clustering_dir / --reference_2_clustering_dir: Used during validation and testing to ensure consistency with the training clusters.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_datasetTo prepare the AlphaFold 3 PDB dataset, you must download both first-assembly and asymmetric unit mmCIF files. It is recommended to use AWS snapshots (e.g., 20240101) for reproducibility.
Warning: The PDB download can require up to 700GB of space.
Steps:
./data/pdb_data/unfiltered_assembly_mmcifs../data/pdb_data/unfiltered_asym_mmcifs..gz files in these directories../data/ccd_data/ and unzip them.# For `assembly1` complexes using AWS snapshot:
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs
# Fallback using 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 using AWS snapshot:
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs
# Fallback using rsync:
rsync -rlpt -v -z --delete --port=33444 \
rsync.rcsb.org::ftp_data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs/
# Download CCD and structural models:
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
# Unzip all downloaded files:
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 {} \;
find data/ccd_data/ -type f -name "*.gz" -exec gzip -d {} \;After downloading the raw PDB files, use the filtering scripts to create the train, validation, and test sets. These scripts process the assembly and asymmetric unit files using the CCD data.
Required arguments:
--mmcif_assembly_dir: Path to first-assembly mmCIF files.--mmcif_asym_dir: Path to asymmetric unit mmCIF files.--ccd_dir: Path to the Chemical Component Dictionary.--output_dir: Desired directory for the filtered mmCIF files (e.g., ./data/pdb_data/train_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>The project includes a Dockerfile configured for PyTorch with GPU support. The default base image is pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime.
Build Options:
PYTORCH_TAG: Change the base image (e.g., to a different PyTorch/CUDA version).GIT_TAG: Change the version of the repository to install.Running the container:
Use the --gpus all flag and mount a local volume to the /data directory to allow the container to access training data.
## Build Docker Container
docker build -t af3 .
## Build with custom 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 af3You can install the package using pip:
$ pip install alphafold3-pytorchA PairwiseBlock is a collection of modules used in both MSAModule and Pairformer. It processes pairwise representations using a combination of triangle modules and transitions.
It consists of:
All modules are wrapped in PreLayerNorm and use residual connections.
forward method signature:
forward(
pairwise_repr: Float['b n n d'],
*,
mask: Bool['b n'] | None = None,
value_residuals: tuple[Tensor, Tensor] | None = None,
return_values = False
)class PairwiseBlock(
Module
):
def __init__(
self,
*,
dim_pairwise = 128,
tri_mult_dim_hidden = None,
tri_attn_dim_head = 32,
tri_attn_heads = 4,
dropout_row_prob = 0.25,
dropout_col_prob = 0.25,
accept_value_residual = False
):
# ... implementation ...
def forward(
self,
pairwise_repr: Float['b n n d'],
*,
mask: Bool['b n'] | None = None,
value_residuals: tuple[Tensor, Tensor] | None = None,
return_values = False
):
# ... implementation ...The Web UI follows a specific workflow for structure prediction:
Molecule type (Protein, DNA, RNA, Ligand, or Ion), specifying the number of Copies, and providing a Sequence or selecting a predefined ligand/ion.Predict button is clicked, the list of entities is passed to the fold function.fold function converts the entities into an Alphafold3Input object and calls model.forward_with_alphafold3_inputs with return_bio_pdb_structures=True..pdb file in the cache directory, and the 3D structure is rendered in the UI using Molecule3D.When preparing data for training, especially with large structures, the project supports windowing atom-pair representations.
If atoms_per_window is provided during collation via collate_inputs_to_batched_atom_input, the following happens:
atompair_inputs are converted from full pairwise representations to windowed versions using full_pairwise_repr_to_windowed.atompair_ids are converted to windowed versions using full_attn_bias_to_windowed.This windowing is crucial for managing memory when dealing with large protein complexes.
The MSAModule implements Algorithm 8 for processing Multiple Sequence Alignments (MSA). It facilitates communication between single representations and pairwise representations through a series of layers. Each layer consists of:
To handle large MSAs, the module can cap the number of MSAs using max_num_msa by sampling without replacement.
dim_single: Dimension of the single representation.dim_pairwise: Dimension of the pairwise representation.depth: Number of MSA layers.dim_msa: Dimension of the MSA features.dim_msa_input: Input dimension for MSA (defaults to NUM_MSA_ONE_HOT).dim_additional_msa_feats: Number of additional MSA features (defaults to 2).max_num_msa: Maximum number of MSAs to process (if exceeded, samples the top $k$ based on random noise).checkpoint: If True, uses gradient checkpointing for the layers to save memory.forward method signature:
forward(
*,
single_repr: Float['b n ds'],
pairwise_repr: Float['b n n dp'],
msa: Float['b s n dm'],
mask: Bool['b n'] | None = None,
msa_mask: Bool['b s'] | None = None,
additional_msa_feats: Float['b s n {self.dmi}'] | None = None
) -> Float['b n n dp']class MSAModule(
Module
):
def __init__(
self,
*,
dim_single = 384,
dim_pairwise = 128,
depth = 4,
dim_msa = 64,
dim_msa_input=NUM_MSA_ONE_HOT,
dim_additional_msa_feats=2,
outer_product_mean_dim_hidden = 32,
msa_pwa_dropout_row_prob = 0.15,
msa_pwa_heads = 8,
msa_pwa_dim_head = 32,
checkpoint = False,
pairwise_block_kwargs: dict = dict(),
max_num_msa: int | None = None,
layerscale_output: bool = True
):
# ... implementation ...
def forward(
self,
*,
single_repr: Float['b n ds'],
pairwise_repr: Float['b n n dp'],
msa: Float['b s n dm'],
mask: Bool['b n'] | None = None,
msa_mask: Bool['b s'] | None = None,
additional_msa_feats: Float['b s n {self.dmi}'] | None = None
) -> Float['b n n dp']:
# ... implementation ...The Trainer can maintain an ema_model to provide a more stable version of the weights for validation and testing. This is controlled by the use_ema flag.
Key Concepts:
ema_update_model_with_ema_every steps) using the ema_decay (beta) value.ema_model if it exists. If use_ema is False, the standard model is used.ema_on_cpu=True) to save GPU memory.The library provides several ways to handle biological data inputs:
Alphafold3Input, AtomInput, MoleculeInput, and PDBInput define the structure of the data.atom_input_to_file and file_to_atom_input for file I/O. Use alphafold3_inputs_to_batched_atom_input or collate_inputs_to_batched_atom_input to prepare data for batch processing.AtomDataset and PDBDataset are available for managing large collections of inputs.register_input_transform allows you to define custom preprocessing steps.