# hannesstark/boltzgen architecture

> Architecture diagram of the GitHub repository hannesstark/boltzgen, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-14). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/hannesstark/boltzgen
- Repository: https://github.com/hannesstark/boltzgen

## Overview

#### 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/`.

Main entry points: [`src/boltzgen/data/tokenize/tokenizer.py`](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/tokenize/tokenizer.py) (Structure tokenizer), [`src/boltzgen/task/task.py`](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/task.py) (Pipeline orchestration), [`src/boltzgen/task/predict/data_from_yaml.py`](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_from_yaml.py) (Design YAML loader), [`src/boltzgen/task/predict/predict.py`](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/predict.py) (Complex refolding)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entry & orchestration"]
  node_cli["BoltzGen CLI<br/>command entry point<br/>[boltzgen.py]"]
  node_task["Pipeline orchestration<br/>task runner<br/>[task.py]"]
end

subgraph group_input["Input preparation"]
  node_yaml_loader["Design YAML loader<br/>specification loader<br/>[data_from_yaml.py]"]
  node_generated_loader["Generated artifact loader"]
  node_ligand_loader["Ligand data loader<br/>molecule input<br/>[data_ligands.py]"]
  node_structure_parser["Structure parser<br/>mmCIF parser<br/>[mmcif.py]"]
  node_tokenizer["Structure tokenizer<br/>token representation<br/>[tokenizer.py]"]
  node_featurizer["Feature pipeline<br/>model features<br/>[featurizer.py]"]
end

subgraph group_model["Design model"]
  node_boltz_model{{"Boltz design model<br/>neural model<br/>[boltz.py]"}}
  node_trunk["Structural trunk<br/>representation encoder<br/>[trunk.py]"]
  node_diffusion["Backbone diffusion<br/>geometry sampler<br/>[diffusion.py]"]
  node_inverse_fold["Inverse folding<br/>sequence designer<br/>[inverse_fold.py]"]
  node_structure_writer["Structure artifacts<br/>CIF/PDB/NPZ output<br/>[mmcif.py]"]
end

subgraph group_evaluate["Validation &amp; selection"]
  node_refold["Complex refolding<br/>prediction stage<br/>[predict.py]"]
  node_affinity["Affinity prediction<br/>small-molecule evaluation<br/>[affinity.py]"]
  node_analysis["Structure analysis<br/>quality metrics<br/>[analyze.py]"]
  node_filter["Filter and rank<br/>selection pipeline<br/>[filter.py]"]
  node_dynamic_filter["Metric threshold rules<br/>filter engine<br/>[filter.py]"]
  node_cluster["Diversity clustering<br/>selection strategy<br/>[cluster.py]"]
  node_final_output["final_ranked_designs<br/>selected design output"]
end

subgraph group_training["Training"]
  node_training_entry["Training entry point<br/>config-driven training<br/>[main.py]"]
end

node_cli -->|"starts"| node_task
node_task -->|"loads design specs"| node_yaml_loader
node_task -.->|"reloads artifacts"| node_generated_loader
node_yaml_loader -->|"resolves ligands"| node_ligand_loader
node_yaml_loader -->|"reads target structures"| node_structure_parser
node_generated_loader -->|"prepares prior designs"| node_tokenizer
node_structure_parser -->|"structures"| node_tokenizer
node_ligand_loader -->|"molecules"| node_tokenizer
node_tokenizer -->|"tokens"| node_featurizer
node_featurizer -->|"structured features"| node_boltz_model
node_boltz_model -->|"encodes"| node_trunk
node_trunk -->|"representations"| node_diffusion
node_diffusion -->|"backbone"| node_inverse_fold
node_inverse_fold -->|"candidate designs"| node_structure_writer
node_structure_writer -->|"candidate artifacts"| node_refold
node_refold -.->|"small-molecule protocol"| node_affinity
node_refold -->|"refolded complex CIFs"| node_analysis
node_affinity -.->|"affinity metrics"| node_analysis
node_analysis -->|"analysis artifacts"| node_filter
node_filter -->|"applies thresholds"| node_dynamic_filter
node_dynamic_filter -->|"qualified designs"| node_cluster
node_cluster -->|"diverse ranked set"| node_final_output
node_training_entry -.->|"trains checkpoints"| node_boltz_model

click node_cli "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/cli/boltzgen.py"
click node_task "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/task.py"
click node_yaml_loader "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_from_yaml.py"
click node_generated_loader "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_from_generated.py"
click node_ligand_loader "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_ligands.py"
click node_structure_parser "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/parse/mmcif.py"
click node_tokenizer "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/tokenize/tokenizer.py"
click node_featurizer "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/feature/featurizer.py"
click node_boltz_model "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/models/boltz.py"
click node_trunk "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/trunk.py"
click node_diffusion "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/diffusion.py"
click node_inverse_fold "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/inverse_fold.py"
click node_structure_writer "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/write/mmcif.py"
click node_refold "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/predict.py"
click node_affinity "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/affinity.py"
click node_analysis "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/analyze/analyze.py"
click node_filter "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/filter/filter.py"
click node_dynamic_filter "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/filter/dynamic/filter.py"
click node_cluster "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/sample/cluster.py"
click node_training_entry "https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/resources/main.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_cli,node_task toneBlue
class node_yaml_loader,node_generated_loader,node_ligand_loader,node_structure_parser,node_tokenizer,node_featurizer toneAmber
class node_boltz_model,node_trunk,node_diffusion,node_inverse_fold,node_structure_writer toneMint
class node_refold,node_affinity,node_analysis,node_filter,node_dynamic_filter,node_cluster,node_final_output toneRose
class node_training_entry toneIndigo
```

## Components

### Entry & orchestration

- [BoltzGen CLI](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/cli/boltzgen.py): command entry point · `src/boltzgen/cli/boltzgen.py`
- [Pipeline orchestration](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/task.py): task runner · `src/boltzgen/task/task.py`

### Input preparation

- [Design YAML loader](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_from_yaml.py): specification loader · Builds prediction input from target and design specifications. · `src/boltzgen/task/predict/data_from_yaml.py`
- [Generated artifact loader](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_from_generated.py): artifact loader · `src/boltzgen/task/predict/data_from_generated.py`
- [Ligand data loader](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/data_ligands.py): molecule input · Resolves ligand and molecular-dictionary data. · `src/boltzgen/task/predict/data_ligands.py`
- [Structure parser](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/parse/mmcif.py): mmCIF parser · `src/boltzgen/data/parse/mmcif.py`
- [Structure tokenizer](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/tokenize/tokenizer.py): token representation · `src/boltzgen/data/tokenize/tokenizer.py`
- [Feature pipeline](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/feature/featurizer.py): model features · Builds token, template, crop, padding, and selection features. · `src/boltzgen/data/feature/featurizer.py`

### Design model

- [Boltz design model](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/models/boltz.py): neural model · Transforms structured features into candidate backbone and sequence artifacts. · `src/boltzgen/model/models/boltz.py`
- [Structural trunk](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/trunk.py): representation encoder · Encodes single and pair structural representations. · `src/boltzgen/model/modules/trunk.py`
- [Backbone diffusion](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/diffusion.py): geometry sampler · `src/boltzgen/model/modules/diffusion.py`
- [Inverse folding](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/inverse_fold.py): sequence designer · Generates sequences with residue-constraint masks. · `src/boltzgen/model/modules/inverse_fold.py`
- [Structure artifacts](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/write/mmcif.py): CIF/PDB/NPZ output · Writes generated and refolded structures with metadata. · `src/boltzgen/data/write/mmcif.py`

### Validation & selection

- [Complex refolding](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/predict/predict.py): prediction stage · Refolds candidates as complexes and optionally isolated binders. · `src/boltzgen/task/predict/predict.py`
- [Affinity prediction](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/model/modules/affinity.py): small-molecule evaluation · Optional protocol stage for small-molecule designs. · `src/boltzgen/model/modules/affinity.py`
- [Structure analysis](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/analyze/analyze.py): quality metrics · Computes structural and interface metrics from refolded complexes. · `src/boltzgen/task/analyze/analyze.py`
- [Filter and rank](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/task/filter/filter.py): selection pipeline · Applies quality/composition rules and ranks survivors. · `src/boltzgen/task/filter/filter.py`
- [Metric threshold rules](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/filter/dynamic/filter.py): filter engine · `src/boltzgen/data/filter/dynamic/filter.py`
- [Diversity clustering](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/data/sample/cluster.py): selection strategy · Selects diverse designs within the requested budget. · `src/boltzgen/data/sample/cluster.py`
- **final_ranked_designs**: selected design output · Selected CIFs, metric CSVs, and overview plots.

### Training

- [Training entry point](https://github.com/hannesstark/boltzgen/blob/main/src/boltzgen/resources/main.py): config-driven training · Runs multi-GPU training with optimizer and EMA support. · `src/boltzgen/resources/main.py`

## Connections

- BoltzGen CLI → Pipeline orchestration: starts
- Pipeline orchestration → Design YAML loader: loads design specs
- Pipeline orchestration → Generated artifact loader: reloads artifacts
- Design YAML loader → Ligand data loader: resolves ligands
- Design YAML loader → Structure parser: reads target structures
- Generated artifact loader → Structure tokenizer: prepares prior designs
- Structure parser → Structure tokenizer: structures
- Ligand data loader → Structure tokenizer: molecules
- Structure tokenizer → Feature pipeline: tokens
- Feature pipeline → Boltz design model: structured features
- Boltz design model → Structural trunk: encodes
- Structural trunk → Backbone diffusion: representations
- Backbone diffusion → Inverse folding: backbone
- Inverse folding → Structure artifacts: candidate designs
- Structure artifacts → Complex refolding: candidate artifacts
- Complex refolding → Affinity prediction: small-molecule protocol
- Complex refolding → Structure analysis: refolded complex CIFs
- Affinity prediction → Structure analysis: affinity metrics
- Structure analysis → Filter and rank: analysis artifacts
- Filter and rank → Metric threshold rules: applies thresholds
- Metric threshold rules → Diversity clustering: qualified designs
- Diversity clustering → final_ranked_designs: diverse ranked set
- Training entry point → Boltz design model: trains checkpoints