boltz

repository·main·Indexed 24 days ago

https://github.com/jwohlwend/boltz

A family of biomolecular interaction prediction models. Boltz-2 is a foundation model that jointly models complex structures and binding affinities to provide high-speed, high-accuracy in silico screening for drug discovery. The package includes a CLI for running inference via `boltz predict`, supporting polymers (protein, DNA, RNA) and ligands, as well as binding affinity prediction for small molecules.

Tokens
17.2K
Snippets
17
Records
75
Agent score
87%

What's inside boltz

  1. Process raw data for training

    main

    If you want to use your own data or replicate the pipeline, follow these steps in the scripts/process directory.

    Prerequisites

    1. Install requirements: pip install -r requirements.txt.
    2. Install external dependencies: mmseqs and redis.

    Step-by-Step Pipeline

    1. Preprocess CCD Dictionary: Use ccd.py with a components.cif file to create ccd.pkl.
    2. Create Sequence Clusters: Use cluster.py with your sequences and mmseqs to assign proteins to similarity clusters.
    3. Generate MSAs: Ensure MSA files are named using the SHA256 hash of the query sequence. For ColabFold-style MSAs, ensure Uniref sequences start with >UniRef100_UNIREFID.
    4. Process MSAs:
      • Start a redis-server using the provided taxonomy.rdb.
      • Run msa.py to annotate sequences with taxonomy IDs.
    5. Process Structures:
      • Start a redis-server using ccd.rdb.
      • Run rcsb.py (for mmcif format) using your clustering JSON to generate processed structural data.
    # 1. Setup environment
    cd scripts/process
    pip install -r requirements.txt
    
    # 2. Preprocess CCD
    python ccd.py --components components.cif --outdir ./ccd
    
    # 3. Clustering
    python cluster.py --ccd ccd.pkl --sequences pdb_seqres.txt --mmseqs PATH_TO_MMSEQS_EXECUTABLE --outdir ./clustering
    
    # 4. MSA Processing
    redis-server --dbfilename taxonomy.rdb --port 7777
    python msa.py --msadir YOUR_MSA_DIR --outdir YOUR_OUTPUT_DIR --redis-port 7777
    
    # 5. Structure Processing
    redis-server --dbfilename ccd.rdb --port 7777
    python rcsb.py --datadir PATH_TO_MMCIF_DIR --cluster clustering/clustering.json --outdir YOUR_OUTPUT_DIR --use-assembly --max-file-size 7000000 --redis-port 7777
  2. Reproduce Boltz-1 evaluation results

    main

    To reproduce Boltz-1 evaluation results, use the following setup:

    Datasets:

    • PDB test set: 541 targets (post-validation cut-off, max 40% protein sequence similarity, 80% ligand Tanimoto similarity).
    • CASP15: 66 difficult targets from the CASP 2022 competition.

    Benchmarking Parameters:

    • MSAs: Use the same MSAs for all models.
    • Recycling steps: 10
    • Sampling steps: 200
    • Samples: 5
    • Metrics: Comparison is based on oracle and top-1 numbers among the 5 samples.

    Software Requirements:

    • Evaluations must be run through OpenStructure version 2.8.0 to ensure reproducibility.
  3. Understand Boltz prediction output structure

    main

    After running a prediction, outputs are organized in an out_dir/ directory. The predictions/ folder contains subdirectories for each input file, containing:

    • [input_file]_model_[N].cif: Predicted structure in CIF format (includes per-token pLDDT scores).
    • confidence_[input_file]_model_0.json: Aggregated confidence scores.
    • affinity_[input_file].json: Predicted affinity scores.
    • pae_[input_file]_model_0.npz: Predicted PAE scores.
    • plddt_[input_file]_model_0.npz: Predicted pLDDT scores.
    • pde_[input_file]_model_0.npz: Predicted PDE scores.

    The processed/ folder contains the intermediate processed data used during inference.

  4. Authenticate to the MSA Server

    main

    When using the --use_msa_server flag, you can authenticate using either Basic Authentication or API Key Authentication. Only one method can be used at a time.

    Basic Authentication

    Use CLI flags or environment variables:

    • CLI: --msa_server_username and --msa_server_password
    • Env Vars: BOLTZ_MSA_USERNAME and BOLTZ_MSA_PASSWORD (recommended)

    API Key Authentication

    Use CLI flags or environment variables:

    • CLI: --api_key_header (default: X-API-Key) and --api_key_value
    • Env Var: MSA_API_KEY_VALUE (recommended)

    Note: If both CLI and environment variables are set, CLI options take precedence.

    # Example: Basic Auth via Env Vars
    export BOLTZ_MSA_USERNAME=myuser
    export BOLTZ_MSA_PASSWORD=mypassword
    boltz predict ... --use_msa_server
    
    # Example: API Key via CLI
    boltz predict ... --use_msa_server --api_key_header X-API-Key --api_key_value <your-api-key>
    
    # Example: API Key via Env Var
    export MSA_API_KEY_VALUE=<your-api-key>
    boltz predict ... --use_msa_server --api_key_header X-API-Key
  5. Access Boltz-1 evaluation files and datasets

    main

    For Boltz-1, evaluation files including MSAs, input files, output files, and evaluation outputs for Boltz-1, Boltz-1x, Chai-1, and AlphaFold3 are available via a Google Drive folder.

    Note: Updated evaluation files for Boltz-2 are coming soon.

    https://drive.google.com/file/d/1JvHlYUMINOaqPTunI9wBYrfYniKgVmxf/view?usp=sharing
  6. Run the Boltz training script

    main

    Execute the training script using scripts/train/train.py and a configuration file.

    Debug Mode

    Before a full run, use the debug=1 flag. This disables DDP (sets to single device), sets num_workers to 0, and disables wandb for easier troubleshooting.

    Standard Training

    Launch the full training run by passing the desired configuration file (e.g., structure.yaml or confidence.yaml).

  7. Install Boltz

    main

    It is recommended to install boltz in a fresh Python environment.

    To install with CUDA support (recommended for GPU users):

    pip install boltz[cuda] -U

    To install for CPU-only or non-CUDA GPU hardware, remove the [cuda] suffix. Note that the CPU version is significantly slower.

    To install directly from GitHub for the latest updates:

    git clone https://github.com/jwohlwend/boltz.git
    cd boltz; pip install -e .[cuda]
    pip install boltz[cuda] -U
  8. Interpret affinity prediction values

    main

    The affinity_[input_file].json file contains two distinct types of predictions that should be used in different contexts:

    1. affinity_probability_binary: Use this for hit-discovery (detecting binders from decoys). It represents the probability (0 to 1) that the ligand is a binder.
    2. affinity_pred_value: Use this for ligand optimization (comparing different active molecules). It reports log10(IC50) derived from μM.
      • Lower values indicate stronger binding.
      • Example: -3 is a strong binder ($10^{-9}$ M), 0 is moderate ($10^{-6}$ M), 2 is a weak binder/decoy ($10^{-4}$ M).

    Conversion to pIC50 (kcal/mol): To convert the model's output y to pIC50 in kcal/mol, use the formula: y --> (6 - y) * 1.364

    {
        "affinity_pred_value": 0.8367,
        "affinity_probability_binary": 0.8425,
        "affinity_pred_value1": 0.8225,
        "affinity_probability_binary1": 0.0,
        "affinity_pred_value2": 0.8225,
        "affinity_probability_binary2": 0.8402
    }
  9. Download pre-processed training datasets

    main

    To run training, download the following pre-processed datasets. Note that approximately 250GB of storage is required for all data.

    RCSB (PDB) Data

    • Structures: rcsb_processed_targets.tar
    • MSAs: rcsb_processed_msa.tar

    OpenFold Data

    • Structures: openfold_processed_targets.tar
    • MSAs: openfold_processed_msa.tar

    Symmetry Data

    • Ligand Symmetry: symmetry.pkl
    # Example: Download and extract RCSB structures
    wget https://boltz1.s3.us-east-2.amazonaws.com/rcsb_processed_targets.tar
    tar -xf rcsb_processed_targets.tar
    rm rcsb_processed_targets.tar
    
    # Example: Download ligand symmetry file
    wget https://boltz1.s3.us-east-2.amazonaws.com/symmetry.pkl
  10. Use structural templates in Boltz YAML input

    main

    The templates field allows you to provide structural templates (CIF or PDB) for protein chains.

    • CIF files: Provide cif: CIF_PATH. You can specify chain_id to map specific chains. Use force: true and a threshold (in Angstroms) to enforce that the backbone does not deviate excessively from the template.
    • PDB files: Provide pdb: PDB_PATH. Boltz incrementally assigns template chain IDs (e.g., A1, A2, B1) based on the PDB file. You can specify chain_id and template_id to explicitly map your input chains to the template chains.
  11. Configure the training process via YAML

    main

    The training script requires a configuration file (templates available in scripts/train/configs/) to specify data paths, output directories, and model parameters.

    Key Parameters for Structure Model Training

    • trainer.devices: Number of devices to use.
    • output: Path to the output directory.
    • resume: Path to a checkpoint file to resume training (set to null if starting fresh).
    • data.datasets: A list of boltz.data.module.training.DatasetConfig objects. Each requires:
      • target_dir: Path to processed structure files.
      • msa_dir: Path to processed MSA files.
    • data.symmetries: Path to the molecule symmetry information file.
    • data.max_tokens: Maximum number of tokens in the input sequence (recommended: 256, 384, or 512).
    • data.max_atoms: Maximum number of atoms in the input structure (recommended: 2304, 3456, or 4608).

    Multi-dataset Configuration Example

    You can define multiple sources (e.g., PDB and OpenFold) by adding multiple entries to the datasets list. Each entry can specify a prob (probability), a sampler (e.g., boltz.data.sample.cluster.ClusterSampler), a cropper (e.g., boltz.data.crop.boltz.BoltzCropper), and a split file.

    trainer:
      devices: 1
    
    output: SET_PATH_HERE
    resume: PATH_TO_CHECKPOINT_FILE
    
    data:
      datasets:
        - _target_: boltz.data.module.training.DatasetConfig
          target_dir: PATH_TO_TARGETS_DIR
          msa_dir: PATH_TO_MSA_DIR
      symmetries: PATH_TO_SYMMETRY_FILE
      max_tokens: 512
      max_atoms: 4608