Protein Structure Prediction with OpenFold

Introduction

AlphaFold2 revolutionized structural biology by predicting 3D protein structures from amino acid sequences with atomic accuracy. However, its original implementation was notoriously difficult to train from scratch and modify. OpenFold is a fast, memory-efficient, and trainable PyTorch reproduction of AlphaFold2. In this tutorial, we will set up OpenFold and run a forward pass to predict the 3D structure of a small protein sequence.

Architecture Overview

graph TD
    Seq(Input Amino Acid Sequence) --> DB[(Protein Sequence Databases)]
    DB -- HHblits / Jackhmmer --> MSA[Multiple Sequence Alignment]
    Seq -- HHsearch --> Temp[Structural Templates]
    
    MSA --> Evo[Evoformer Stack: Extracts Co-evolutionary Features]
    Temp --> Evo
    
    Evo --> Struct[Structure Module: 3D Equivariant Transformer]
    Struct --> PDB(Output: 3D PDB Coordinates + pLDDT Confidence)
        

Prerequisites

• Linux Environment (Ubuntu 20.04+ recommended)
• A GPU with at least 16GB VRAM (e.g., A100, RTX 3090, A4000)
• Conda installed
• Familiarity with FASTA files and PDB formats

Step 1: Environment Setup

OpenFold requires a highly specific environment due to its complex custom CUDA kernels and reliance on DeepMind's original OpenMM dependencies.

git clone https://github.com/aqlaboratory/openfold.git
cd openfold

# Create the conda environment
conda env create -f environment.yml
conda activate openfold_env

# Install OpenFold module and its custom CUDA kernels
python setup.py install

Step 2: Downloading Model Parameters

OpenFold can use either the original AlphaFold2 weights or its own OpenFold weights (which were trained from scratch and perform equally well, if not better on some tasks). We'll download the OpenFold weights.

# Run the download script provided by the repository
scripts/download_openfold_params.sh openfold/resources

You will also need sequence databases (UniRef90, MGnify, BFD, PDB70) to generate Multiple Sequence Alignments (MSAs). These databases are massive (~2TB uncompressed). For testing, you can use the reduced databases provided by AlphaFold, or mock MSAs.

Step 3: Preparing the Input Sequence

Create a FASTA file containing the amino acid sequence of the protein you want to fold.

cat << EOF > input.fasta
>test_protein
MQYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE
EOF

Step 4: Generating the Multiple Sequence Alignment (MSA) using MMseqs2

AlphaFold/OpenFold's accuracy comes from evolutionary information encoded in MSAs. Traditionally, generating these alignments against local 2TB databases using Jackhmmer takes hours. As an alternative, we can use the ColabFold MMseqs2 API to fetch MSAs in seconds over the network.

import requests
import time
import os

sequence = "MQYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE"
job_id = requests.post("https://api.colabfold.com/ticket/msa", data={"q": f">test\n{sequence}"}).json()["id"]

print("Waiting for MSA...")
while True:
    status = requests.get(f"https://api.colabfold.com/ticket/{job_id}").json()["status"]
    if status == "COMPLETE":
        break
    time.sleep(5)

# Download the a3m MSA file
os.makedirs("alignments", exist_ok=True)
msa_data = requests.get(f"https://api.colabfold.com/result/download/{job_id}").content
with open("alignments/test_protein.a3m", "wb") as f:
    f.write(msa_data)
print("MSA successfully downloaded to alignments/test_protein.a3m")

Step 5: Running Inference

Now that we have the sequence, the MSAs, and the model weights, we can run the forward pass to predict the 3D coordinates of every atom in the protein. Because we precomputed the MSA via API, this step will be extremely fast.

python run_pretrained_openfold.py \
    input.fasta \
    openfold/resources/openfold_params/finetuning_ptm_2.pt \
    --output_dir ./predictions \
    --alignment_dir ./alignments \
    --use_precomputed_alignments

Step 6: Analyzing the Output

The script will output several files in the ./predictions directory, most notably a .pdb file. This file contains the 3D coordinates (X, Y, Z) of every atom. The B-factor column in this PDB file actually contains the pLDDT score (Predicted Local Distance Difference Test)—a confidence metric between 0 and 100 for each amino acid residue.

You can visualize this .pdb file using PyMOL or ChimeraX, coloring the structure by the B-factor column to instantly see which parts of the prediction the model is highly confident in (usually > 90) versus unstructured or uncertain regions (usually < 50).

Conclusion

OpenFold provides an accessible, modifiable PyTorch implementation of the AlphaFold architecture. Because it is highly optimized, it allows researchers to not just run inference, but actually fine-tune the model on proprietary structural datasets, modify the neural network layers, or train entirely new models from scratch on specialized hardware clusters. Bypassing the local Jackhmmer searches in favor of MMseqs2 API calls dramatically reduces the barrier to entry for independent researchers.