TEMPO-Flow: conditional flow matching of subcellular RNA dynamics

Supporting code for the TEMPO-Flow paper

Table of contents

• About this repository
• Quick start
  • Installation
  • Data preparation
• Training
• Evaluation
• Codebase organization
• License

About this repository

TEMPO-Flow learns a single gene-token transformer velocity field $v_\theta(x, t, c)$ over the nuclear and cytoplasmic RNA compartments. Integrating the field forecasts held-out pulse–chase distributions, and a sparse-physics head reads out interpretable transcription / degradation rate constants ($k_\mathrm{tx}$, $k_\mathrm{deg}$) on held-out genes.

Everything is built on PyTorch Lightning + Hydra — every model is a LightningModule, every dataset a LightningDataModule, configured through configs/. The default configuration is the champion model (the E_large gene-token compartment flow).

Primary dataset: TEMPOmap pulse–chase 3D in-situ sequencing (Ren et al., Nat. Methods 20, 695–705, 2023).

Quick start

Installation

conda env create -f environment.yml      # creates the `tempo-flow` env
conda activate tempo-flow
pip install -e .                          # installs the `tempo_flow` package

(Or pip install -r requirements.txt && pip install -e . in an existing Python ≥ 3.11 environment.)

Data preparation

The TEMPOmap pulse–chase data is not redistributed here — obtain it from the original study (Ren et al., Nat. Methods, 2023). Then build the model-ready inputs:

bash   scripts/download_all.sh               # fetch public companion data
python scripts/tempomap_csv_to_anndata.py    # raw TEMPOmap CSV -> AnnData
python scripts/make_splits.py                # train / val / test split
python scripts/make_5fold_gene_splits.py     # 5-fold gene cross-validation
python scripts/build_gene_embeddings.py      # gene-token embeddings
python scripts/verify_data.py                # sanity checks

See each script's header / --help for paths and options.

Training

Hydra entrypoint — the defaults are the champion model:

python scripts/train.py                       # champion (E_large compartment)
bash   scripts/run_champion_seeds.sh          # multi-seed champion

Override any config group from the command line, e.g. python scripts/train.py data.batch_size=64 seed=2025. Optional Weights & Biases logging activates automatically if WANDB_API_KEY is set.

Evaluation

python scripts/evaluate.py                   # held-out forecasting metrics
python scripts/eval_rates_multi_v2.py        # rate recovery (k_tx, k_deg)
python scripts/eval_rates_external.py        # cross-assay rate validation

Codebase organization

src/tempo_flow/
  models/             gene-token velocity transformer, sparse-physics rate head
  lightning_modules/  the TEMPO-Flow training module (CFM + endpoint + flux + rate)
  data/               TEMPOmap LightningDataModule
  training/           callbacks (EMA, progress logging)
  eval/               evaluation utilities
configs/              Hydra configs: model/, data/, train/, eval/
scripts/              data preparation, training, and evaluation

License

This project is released under the MIT License — see LICENSE.