hannesstark/boltzgen

Purpose and Entry Points

BoltzGen is a Python/GPU molecular-binder design system. It generates candidate protein or peptide backbones, designs sequences, refolds and evaluates them, then selects a quality- and diversity-ranked output set. User-facing commands are implemented in src/boltzgen/cli/boltzgen.py; pipeline orchestration starts in src/boltzgen/task/task.py. Training has a separate config-driven entry point at src/boltzgen/resources/main.py.

Specification-to-Model Input

Design YAML files define target structures, designed chains, ligands, binding regions, structural visibility, and bonds. Prediction loaders convert those specifications or prior generated artifacts into model-ready data:

• src/boltzgen/task/predict/data_from_yaml.py

• src/boltzgen/task/predict/data_from_generated.py

• src/boltzgen/task/predict/data_ligands.py

The input boundary accepts local PDB/mmCIF files and YAML-relative paths; molecule definitions may come from a downloaded molecular dictionary.

Structural Data Pipeline

The data layer parses structures and produces token-, feature-, and template-based representations consumed by models. It also handles cropping, padding, selection, and structure serialization:

• src/boltzgen/data/parse/mmcif.py

• src/boltzgen/data/tokenize/tokenizer.py

• src/boltzgen/data/feature/featurizer.py

Generated and refolded structures are emitted as mmCIF/PDB plus NPZ metadata through src/boltzgen/data/write/mmcif.py and src/boltzgen/data/write/pdb.py.

Generative Design and Inverse Folding

The design stage uses a Boltz-based neural architecture: a trunk encodes single/pair structural representations and a diffusion module samples designed backbone geometry. Sequence generation is then handled by inverse folding, including residue-constraint masks.

• src/boltzgen/model/models/boltz.py

• src/boltzgen/model/modules/trunk.py

• src/boltzgen/model/modules/diffusion.py

• src/boltzgen/model/modules/inverse_fold.py

This is the central model boundary: structured features enter the model; candidate backbone/sequence artifacts leave it for downstream validation.

Refolding, Affinity, and Analysis

Candidate designs are refolded as complexes, and optionally as isolated binders; small-molecule protocols additionally invoke affinity prediction. Analysis computes structural and interface quality metrics over refolded outputs.

• src/boltzgen/task/predict/predict.py

• src/boltzgen/task/analyze/analyze.py

• src/boltzgen/model/modules/affinity.py

Protocol defaults determine whether isolated design folding and affinity are run. Refolded complex CIFs are the primary analysis/filtering input.

Filtering and Final Selection

Filtering applies metric thresholds and composition/quality rules, ranks survivors, then performs diversity-aware selection within the requested budget. It can be rerun independently against existing analysis artifacts.

• src/boltzgen/task/filter/filter.py

• src/boltzgen/data/filter/dynamic/filter.py

• src/boltzgen/data/sample/cluster.py

The resulting output boundary is final_ranked_designs, containing selected CIFs, metric CSVs, and overview plots.

Runtime, Configuration, and External Dependencies

Python 3.11+ runs the package; GPU execution is expected for practical design workloads, with Docker and SLURM deployment paths (Dockerfile, slurm-example/run_job_array.slurm). Step and protocol configuration is resource-backed:

• src/boltzgen/resources/config/design.yaml

• src/boltzgen/resources/config/fold.yaml

• src/boltzgen/resources/config/filtering.yaml

Model checkpoints, inference molecule data, and training datasets are downloaded from Hugging Face and cached locally. Training uses multi-GPU configurations and includes optimizer/EMA support in src/boltzgen/model/optim/.