# lotfollahi-lab/perturbgen architecture

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

- Interactive diagram: https://gitdiagram.com/lotfollahi-lab/perturbgen
- Repository: https://github.com/lotfollahi-lab/perturbgen

## Overview

PerturbGen is a single-cell generative model for predicting perturbation effects along cellular trajectories. Its principal workflow curates and tokenizes transcriptomes, trains a masking encoder and trajectory/count model, then applies in-silico perturbations to predict downstream states. The README also describes cell/gene embeddings, context-specific gene-program discovery, and perturbation atlases (PIPs). These are represented as product capabilities, but their detailed implementation paths are not established by the excerpts. The SOURCE INDEX supports the main training, model, and data-module relationships; unsampled analysis links are documented at capability level only.

Main entry points: [`perturbgen/train.py`](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/train.py) (Trajectory training entry), [`perturbgen/Perturb/trainer.py`](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/trainer.py) (Perturbation engine), [`perturbgen/Perturb/val.py`](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/val.py) (Perturbation inference), [`scmaskgit/train.py`](https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/train.py) (Masking training entry)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_prep["Cohort preparation"]
  node_cohort[("Transcriptome cohort")]
  node_harmonize["Cohort harmonization"]
  node_cohort_tokenize["Cohort tokenization<br/>[tokenization.py]"]
  node_geneformer_tokenize["Gene tokenization<br/>[GF_tokenisation.py]"]
  node_tokenizer["Token utilities<br/>[tokenizer.py]"]
end

subgraph group_pretrain["Masking pretraining"]
  node_mask_data["Masking data loader<br/>[datamodule.py]"]
  node_mask_model["Masking transformer<br/>[T_model.py]"]
  node_mask_trainer["Masking trainer<br/>[trainer.py]"]
  node_mask_train["Masking training entry<br/>[train.py]"]
end

subgraph group_trajectory["Trajectory modeling"]
  node_trajectory_data["Trajectory data loader<br/>[datamodule.py]"]
  node_trajectory_model["Trajectory transformer<br/>[transformer.py]"]
  node_trajectory_trainer["Trajectory trainer<br/>[trainer.py]"]
  node_training_entry["Trajectory training entry<br/>[train.py]"]
  node_api["Python model API<br/>[model.py]"]
end

subgraph group_perturb["Perturbation inference"]
  node_perturb_trainer["Perturbation engine<br/>[trainer.py]"]
  node_perturb_model["Perturbation transformer<br/>[transformer.py]"]
  node_perturb_validation["Perturbation inference<br/>[val.py]"]
  node_predictions["Predicted cell states<br/>[utils.py]"]
end

subgraph group_analysis["Biological analysis"]
  node_embedding_analysis["Embedding analysis<br/>[model.py]"]
  node_program_discovery["Gene program discovery"]
  node_perturbation_atlas["PIP atlas analysis"]
  node_metrics["Prediction metrics<br/>[metric.py]"]
end

node_researcher(("Researcher"))
node_wandb(("Weights &amp; Biases"))

node_researcher -->|"provides data"| node_cohort
node_cohort -->|"curated by"| node_harmonize
node_harmonize -->|"precedes tokenization"| node_cohort_tokenize
node_cohort_tokenize -.->|"tokenizes cells"| node_geneformer_tokenize
node_geneformer_tokenize -->|"uses"| node_tokenizer
node_researcher -->|"starts training"| node_mask_train
node_mask_train -->|"loads batches"| node_mask_data
node_mask_train -->|"configures training"| node_mask_trainer
node_mask_trainer -->|"trains"| node_mask_model
node_mask_model -->|"encoder reference"| node_trajectory_model
node_researcher -->|"starts training"| node_training_entry
node_training_entry -->|"loads batches"| node_trajectory_data
node_training_entry -->|"configures training"| node_trajectory_trainer
node_trajectory_trainer -->|"trains"| node_trajectory_model
node_researcher -->|"uses model API"| node_api
node_api -->|"invokes training"| node_training_entry
node_api -->|"invokes inference"| node_perturb_validation
node_perturb_validation -->|"loads input data"| node_trajectory_data
node_perturb_validation -->|"configures inference"| node_perturb_trainer
node_perturb_trainer -->|"inherits count decoder"| node_trajectory_trainer
node_perturb_trainer -->|"applies intervention"| node_perturb_model
node_perturb_trainer -->|"writes predictions"| node_predictions
node_predictions -.->|"evaluated against data"| node_metrics
node_researcher -->|"requests embeddings"| node_embedding_analysis
node_embedding_analysis -->|"supports discovery"| node_program_discovery
node_predictions -->|"aggregated into PIPs"| node_perturbation_atlas
node_researcher -->|"analyzes perturbations"| node_perturbation_atlas
node_training_entry -.->|"logs runs"| node_wandb
node_perturb_validation -.->|"logs inference"| node_wandb

click node_cohort "https://github.com/lotfollahi-lab/perturbgen/tree/main/pretraining_cohort"
click node_harmonize "https://github.com/lotfollahi-lab/perturbgen/blob/main/pretraining_cohort/preprocessing/harmonize_bronze_to_silver.py"
click node_cohort_tokenize "https://github.com/lotfollahi-lab/perturbgen/blob/main/pretraining_cohort/tokenization/tokenization.py"
click node_geneformer_tokenize "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/pp/GF_tokenisation.py"
click node_tokenizer "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/pp/tokenizer.py"
click node_mask_data "https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/Dataloaders/datamodule.py"
click node_mask_model "https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/Modules/T_model.py"
click node_mask_trainer "https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/Model/trainer.py"
click node_mask_train "https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/train.py"
click node_trajectory_data "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Dataloaders/datamodule.py"
click node_trajectory_model "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Modules/transformer.py"
click node_trajectory_trainer "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Model/trainer.py"
click node_training_entry "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/train.py"
click node_api "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/model.py"
click node_perturb_trainer "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/trainer.py"
click node_perturb_model "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/transformer.py"
click node_perturb_validation "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/val.py"
click node_predictions "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/src/utils.py"
click node_embedding_analysis "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/model.py"
click node_metrics "https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/src/metric.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_cohort,node_harmonize,node_cohort_tokenize,node_geneformer_tokenize,node_tokenizer toneBlue
class node_mask_data,node_mask_model,node_mask_trainer,node_mask_train toneAmber
class node_trajectory_data,node_trajectory_model,node_trajectory_trainer,node_training_entry,node_api toneMint
class node_perturb_trainer,node_perturb_model,node_perturb_validation,node_predictions toneRose
class node_embedding_analysis,node_program_discovery,node_perturbation_atlas,node_metrics,node_researcher,node_wandb toneIndigo
```

## Components

### Cohort preparation

- [Transcriptome cohort](https://github.com/lotfollahi-lab/perturbgen/tree/main/pretraining_cohort): `pretraining_cohort`
- [Cohort harmonization](https://github.com/lotfollahi-lab/perturbgen/blob/main/pretraining_cohort/preprocessing/harmonize_bronze_to_silver.py): `pretraining_cohort/preprocessing/harmonize_bronze_to_silver.py`
- [Cohort tokenization](https://github.com/lotfollahi-lab/perturbgen/blob/main/pretraining_cohort/tokenization/tokenization.py): `pretraining_cohort/tokenization/tokenization.py`
- [Gene tokenization](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/pp/GF_tokenisation.py): `perturbgen/pp/GF_tokenisation.py`
- [Token utilities](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/pp/tokenizer.py): `perturbgen/pp/tokenizer.py`

### Masking pretraining

- [Masking data loader](https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/Dataloaders/datamodule.py): `scmaskgit/Dataloaders/datamodule.py`
- [Masking transformer](https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/Modules/T_model.py): `scmaskgit/Modules/T_model.py`
- [Masking trainer](https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/Model/trainer.py): `scmaskgit/Model/trainer.py`
- [Masking training entry](https://github.com/lotfollahi-lab/perturbgen/blob/main/scmaskgit/train.py): `scmaskgit/train.py`

### Trajectory modeling

- [Trajectory data loader](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Dataloaders/datamodule.py): `perturbgen/Dataloaders/datamodule.py`
- [Trajectory transformer](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Modules/transformer.py): `perturbgen/Modules/transformer.py`
- [Trajectory trainer](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Model/trainer.py): `perturbgen/Model/trainer.py`
- [Trajectory training entry](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/train.py): `perturbgen/train.py`
- [Python model API](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/model.py): `perturbgen/model.py`

### Perturbation inference

- [Perturbation engine](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/trainer.py): `perturbgen/Perturb/trainer.py`
- [Perturbation transformer](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/transformer.py): `perturbgen/Perturb/transformer.py`
- [Perturbation inference](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/Perturb/val.py): `perturbgen/Perturb/val.py`
- [Predicted cell states](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/src/utils.py): `perturbgen/src/utils.py`

### Biological analysis

- [Embedding analysis](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/model.py): `perturbgen/model.py`
- **Gene program discovery**
- **PIP atlas analysis**
- [Prediction metrics](https://github.com/lotfollahi-lab/perturbgen/blob/main/perturbgen/src/metric.py): `perturbgen/src/metric.py`

### Other components

- **Researcher**
- **Weights & Biases**

## Connections

- Researcher → Transcriptome cohort: provides data
- Transcriptome cohort → Cohort harmonization: curated by
- Cohort harmonization → Cohort tokenization: precedes tokenization
- Cohort tokenization → Gene tokenization: tokenizes cells
- Gene tokenization → Token utilities: uses (evidence: `perturbgen/pp/GF_tokenisation.py`)
- Researcher → Masking training entry: starts training
- Masking training entry → Masking data loader: loads batches (evidence: `scmaskgit/train.py`)
- Masking training entry → Masking trainer: configures training (evidence: `scmaskgit/train.py`)
- Masking trainer → Masking transformer: trains (evidence: `scmaskgit/Model/trainer.py`)
- Masking transformer → Trajectory transformer: encoder reference (evidence: `perturbgen/Modules/transformer.py`)
- Researcher → Trajectory training entry: starts training
- Trajectory training entry → Trajectory data loader: loads batches (evidence: `perturbgen/train.py`)
- Trajectory training entry → Trajectory trainer: configures training (evidence: `perturbgen/train.py`)
- Trajectory trainer → Trajectory transformer: trains (evidence: `perturbgen/Model/trainer.py`)
- Researcher → Python model API: uses model API (evidence: `perturbgen/model.py`)
- Python model API → Trajectory training entry: invokes training (evidence: `perturbgen/model.py`)
- Python model API → Perturbation inference: invokes inference (evidence: `perturbgen/model.py`)
- Perturbation inference → Trajectory data loader: loads input data (evidence: `perturbgen/Perturb/val.py`)
- Perturbation inference → Perturbation engine: configures inference (evidence: `perturbgen/Perturb/val.py`)
- Perturbation engine → Trajectory trainer: inherits count decoder (evidence: `perturbgen/Perturb/trainer.py`)
- Perturbation engine → Perturbation transformer: applies intervention (evidence: `perturbgen/Perturb/trainer.py`)
- Perturbation engine → Predicted cell states: writes predictions (evidence: `perturbgen/Perturb/trainer.py`)
- Predicted cell states → Prediction metrics: evaluated against data
- Researcher → Embedding analysis: requests embeddings
- Embedding analysis → Gene program discovery: supports discovery
- Predicted cell states → PIP atlas analysis: aggregated into PIPs
- Researcher → PIP atlas analysis: analyzes perturbations
- Trajectory training entry → Weights & Biases: logs runs
- Perturbation inference → Weights & Biases: logs inference