OpenFold3

repository·main·Indexed 21 days ago

https://github.com/aqlaboratory/openfold-3

A biomolecular structure prediction model designed as a bitwise reproduction of DeepMind's AlphaFold 3. OpenFold3 supports predicting structures for proteins, RNA, DNA (standard and non-canonical), and small molecules. It features multi-modal prediction, MSA pipelines via ColabFold or JackHMMER/hhblits, and acceleration through cuEquivariance and DeepSpeed4Science kernels. The project includes support for distributed predictions across multiple GPUs and resource management for memory-constrained environments.

Tokens
31.5K
Snippets
96
Records
134
Agent score
72%

What's inside openfold3

  1. Overview of OpenFold3-preview features

    main

    OpenFold3-preview is a research preview aiming for bitwise reproduction of AlphaFold3. Key features include:

    • Multi-modal Prediction: Structure prediction for protein, RNA, DNA (standard and non-canonical), and small molecules.
    • MSA Pipelines: Support for generating Multiple Sequence Alignments (MSAs) via the ColabFold server or using JackHMMER / hhblits following the AlphaFold3 protocol.
    • Acceleration: Kernel acceleration via cuEquivariance and DeepSpeed4Science kernels.
    • Scalability: Support for multi-query jobs with distributed predictions across multiple GPUs.
    • Resource Management: Custom settings for memory-constrained GPU environments.
    • Templates: Support for protein monomer structure templates.
  2. Overview of OpenFold3 Inference features

    main

    OpenFold3 provides structure prediction capabilities for various molecule types. While all features described in the AlphaFold3 Nature paper are supported in training, inference support is being rolled out incrementally.

    Supported Molecule Types

    Protein

    • MSA-based prediction: Supports the ColabFold MSA pipeline, pre-computed MSAs, and OpenFold3's native MSA generation pipeline.
    • MSA-free prediction: Prediction without Multiple Sequence Alignments.
    • Template-based prediction: Supports ColabFold template alignments, pre-computed template alignments, and direct CIF template files (no alignments required).
    • Non-canonical residues.

    DNA

    • MSA-free prediction: Prediction without MSA (following AF3 defaults).
    • Non-canonical residues.

    RNA

    • MSA-based prediction: Uses OpenFold3's native MSA generation pipeline.
    • MSA-free prediction.
    • Non-canonical residues.

    Ligand

    • Non-covalent ligands.

    Coming Soon

    • Covalently modified residues and cross-chain covalent bonds (Protein, DNA, RNA).
    • Template-based prediction for RNA.
    • Protein-RNA MSA pairing.
    • Covalently bound ligands and polymeric ligands (e.g., glycans).
  3. Features of OpenFold3-preview

    main

    OpenFold3-preview is an open-source (Apache 2) biological structure prediction model based on AlphaFold3. Key capabilities include:

    • Broad Molecule Support: Predicts structures for standard and non-canonical protein, RNA, DNA chains, and small molecules.
    • MSA Pipelines: Supports generating Multiple Sequence Alignments (MSAs) using the ColabFold server or via JackHMMER / hhblits following the AlphaFold3 protocol.
    • Protein Templates: Supports structure templates for protein monomers.
    • Performance Acceleration: Utilizes kernel acceleration through cuEquivariance and DeepSpeed4Science kernels.
    • Scalability: Supports multi-query jobs and distributed predictions across multiple GPUs.
    • Resource Management: Includes custom settings for memory-constrained GPU resources.
  4. Understand processed MSA outputs

    main

    When --use-msa-server=True is enabled, OpenFold3 generates processed MSAs as .npz files. These are stored in a dedicated MSA output directory.

    Directory Structure

    • main/: Contains .npz files for each unique chain (unpaired MSAs).
    • paired/: Contains .npz files for heteromeric protein complexes (paired MSAs). The filename is a hash of the concatenated sequences.
    • mappings/: Contains metadata to map hashes back to sequences and chain IDs:
      • chain_id_to_rep_id.json
      • query_name_to_complex_id.json
      • rep_id_to_seq.json (hash to sequence)
      • seq_to_rep_id.json (sequence to hash)
    • raw/: Contains the original ColabFold .a3m files. Note: This directory is deleted by default after processing to save space. To keep them, you must modify the source code in colabfold_msa_server.py.
    <msa_output_directory>
    ├── main
    │   ├── <hash of sequence A>.npz
    │   └── <hash of sequence B>.npz
    ├── mappings
    │   ├── chain_id_to_rep_id.json
    │   ├── query_name_to_complex_id.json
    │   ├── README.md
    │   ├── rep_id_to_seq.json
    │   └── seq_to_rep_id.json
    └── paired
        └── <hash of concatenation of sequences A and B>
            ├── <hash of sequence A>.npz
            └── <hash of sequence B>.npz
  5. Customize Inference Settings with runner.yml

    main

    The runner.yml file allows for extensive customization of model and data parameters. Command-line arguments take precedence over settings in the YAML file.

    Common customization scenarios include:

    • Distributed Inference: Use pl_trainer_args to run on multiple GPUs or nodes.
    • Output Formats: Change structure output from cif to pdb via output_writer_settings.
    • Model Updates: Use the model_update field to apply presets (like low_mem) or custom settings.
  6. Understand the SDF File Format for Reference Molecules

    main

    Reference molecules are stored as SDF files in the reference_molecules/ folder using the V2000 molfile format. They include two custom annotation fields for integration with AtomArrays:

    1. atom_annot_atom_name: Canonical atom names for the ligand.
    2. atom_annot_used_atom_mask: A per-atom mask. True for all atoms if a conformer was successfully generated. If CCD-deposited fallback coordinates were used, atoms with missing coordinates are marked False.
         RDKit          3D
    
      4  3  0  0  0  0  0  0  0  0999 V2000
        0.8140   -0.0385    0.0035 C   0  0  0  0  0  0  0  0  0  0  0  0
        1.5232   -1.0718    0.0134 O   0  0  0  0  0  0  0  0  0  0  0  0
        1.5226    1.1766   -0.0223 O   0  0  0  0  0  1  0  0  0  0  0  0
       -0.6545   -0.0200    0.0167 C   0  0  0  0  0  0  0  0  0  0  0  0
      1  2  2  0
      1  3  1  0
      1  4  1  0
    M  CHG  1   3  -1
    M  END
    
    >  <atom_annot_atom_name>  (1)
    C O OXT CH3
    
    >  <atom_annot_used_atom_mask>  (1)
    True True True True
    
    $$$$
  7. How OpenFold3 template modes work

    main

    OpenFold3 supports two distinct modes for using template information during inference. Choosing between them depends on whether you have precomputed alignments or only have structure files.

    1. Alignment-Based Mode (Traditional): This mode is used when you have existing query-to-template alignments. It requires two components:

      • Template Alignments: Files describing how the query sequence relates to the template sequences.
      • Template Structures: The actual structural files (currently only .cif format supported).
    2. CIF Direct Mode (Simplified): This mode is used when you only have template structure files. The system automatically handles the alignment of template chains to your query sequence and selects the best match. You only need to provide the .cif files.

    Note: The template pipeline currently supports monomeric templates and is tested for protein chains only.

  8. High-level structure of OpenFold3 input JSON

    main

    The OpenFold3 inference pipeline uses a single JSON file to specify prediction targets and options. The top-level structure must contain a queries dictionary. Each key in this dictionary represents a unique prediction target (a bioassembly) and is used to name the resulting output files or directories. To perform batch inference, include multiple query entries within this dictionary.

    Note: Do not include a seeds field in your input JSON. Seeds should be managed via the --num-model-seeds CLI argument or within the runner.yml configuration.

    {
      "queries": {
        "query_1": { ... },
        "query_2": { ... }
      }
    }
  9. How OpenFold3 resolves checkpoint paths

    main

    OpenFold3 follows a specific hierarchy to locate parameter files. You can control this via the following parameters:

    1. inference-ckpt-path: If provided, OpenFold3 uses this exact file path immediately.
    2. inference-ckpt-name: If inference-ckpt-path is not provided, OpenFold3 looks for this specific checkpoint name.
    3. Default: If neither is provided, OpenFold3 defaults to openfold3-p2-155k.

    Download and Cache Logic: If the checkpoint is not found locally, OpenFold3 will attempt to download it. The download location depends on your environment configuration:

    • If $OPENFOLD_CACHE is set or cache-path is specified: It checks for a ckpt_root within the cache. If found, it downloads to $OPENFOLD_CACHE/ckpt_root; otherwise, it downloads directly to $OPENFOLD_CACHE.
    • If no cache is configured: It downloads parameters to the default directory ~/.openfold3.
  10. AI contribution guidelines for OpenFold

    main

    OpenFold allows the use of LLMs and coding agents, but contributors must adhere to the following quality and engagement standards to conserve core team resources:

    • Human Responsibility: All contributions (AI-generated or otherwise) are the sole responsibility of the human contributor. You must fully understand, review, and verify that the code works as intended before submission. Always provide test examples demonstrating functionality.
    • Human-Written Communication: Issues and Pull Requests must be written by humans to facilitate meaningful dialogue with the development team. Exceptions are permitted for translations or generating failing test examples, provided the use of AI tools is disclosed.
    • Good First Issues: These issues are reserved for new human developers to learn the workflow. Agentic/AI contributions on issues labeled good first issue will be closed.

    Failure to meet these standards may result in the closure of contributions.

  11. Understand the Validation Cache structure

    main

    The validation cache (validation_cache.json) is used for proper model evaluation. It follows a similar structure to the training cache but includes additional fields for homology and quality assessment.

    To create this cache, you must use the scripts/data_preprocessing/create_pdb_validation_dataset_cache.py script. Note that this process requires the training cache as input to perform homology comparisons.

    Key characteristics of the validation cache include:

    • Stricter filtering: Requires resolution $\le$ 4.5 Angstrom and adheres to specific token limits.
    • Homology detection: Uses a 40% sequence identity threshold for proteins and a 0.85 Tanimoto similarity threshold for ligands.
    • Subsets: Contains both multimer and monomer validation sets.
    "7vl5": {
        "release_date": "2022-03-09",
        "resolution": 1.93,
        "token_count": 1463,
        "chains": {
            "1": {
                "label_asym_id": "A",
                "auth_asym_id": "A",
                "entity_id": 1,
                "molecule_type": "PROTEIN",
                "reference_mol_id": null,
                "alignment_representative_id": "7vl6_A",
                "template_ids": [
                    "5gsl_A",
                    "5gsl_B",
                    "6jow_A"
                ],
                "cluster_id": "2024",
                "cluster_size": null,
                "low_homology": true,
                "metric_eligible": true,
                "use_metrics": false,
                "ranking_model_fit": null,
                "source_subset": "base"
            },
            "3": {
                "label_asym_id": "C",
                "auth_asym_id": "A",
                "entity_id": 2,
                "molecule_type": "LIGAND",
                "reference_mol_id": "BOG",
                "alignment_representative_id": null,
                "template_ids": null,
                "cluster_id": "5086",
                "cluster_size": null,
                "low_homology": false,
                "metric_eligible": false,
                "use_metrics": false,
                "ranking_model_fit": 0.9312,
                "source_subset": "base"
            }
        },
        "interfaces": {
            "1_3": {
                "cluster_id": "2024_5086",
                "cluster_size": null,
                "low_homology": true,
                "metric_eligible": true,
                "use_metrics": false,
                "source_subset": "base"
            }
        }
    }
  12. Understand the OF3 MSA processing pipeline

    main

    The OpenFold3 (OF3) MSA pipeline processes per-chain alignments from multiple sequence databases (e.g., UniRef90, UniProt, ColabFold DB, MGnify) through several stages:

    1. Vertical Concatenation: Alignments from different databases are stacked vertically to form a single MSA block (capped at 16k sequences).
    2. Cross-chain Pairing: Sequences originating from the same species are paired across different chains in a complex to improve protein-protein interface prediction.
    3. Sampling: The combined MSA is sampled down to a target size defined by $[L_i, nJ]$, where $L_i$ is the sequence length and $nJ$ is the number of sampled rows.
    4. Subsampling: During the model forward pass, the MSAModuleEmbedder further subsamples 1024 random rows per recycle, controlled by the max_subsampled_all_msa parameter.