language: en
license: mit
tags:
- biology
- single-cell
- perturbation-prediction
- flow-matching
- pytorch
PIVOT
Perturbation-Informed Vector-Field Optimization for Transcriptomic State Control
With PIVOT, we predict the expression response to a genetic intervention and nominate genes or combinations for a desired response. To this end, we train a perturbation-conditioned endpoint map and use its predictions to rank candidates or optimize intervention embeddings. We then project each optimized embedding to admissible interventions and rerank them by predicted endpoint reward.
The initial Norman experiments recovered 20 of 26 held-out gene pairs within five nominations. A population loss reduced the recorded MMD from 0.3034 to 0.0507 while retaining effect correlation above 0.83. The manuscript and result provenance describe those experiments, their preprocessing scope, and the independent-outcome evaluation.
Install and run
Python 3.11 to 3.13 is supported. The small example runs on a CPU. CUDA is optional for full training.
python -m pip install -e '.[test]'
python scripts/quickstart.py --output runs/example
python -m pytest -q
The example starts from 920 measured Norman cells across 20 noncontrol conditions. It fits training-only expression features and PCA, trains for two epochs, evaluates PIVOT and ridge regression, predicts responses, runs three nomination procedures, and exports a comparison plot. It writes the cache, checkpoints, per-condition results, candidate lists, and plots under runs/example/. Use a new output directory for another run.
The short run verifies the complete workflow. Full benchmark reproduction requires the original datasets and the configurations described below.
Full datasets and training
The adapters use scPerturb Zenodo record 10044268, with fixed filenames and checksum verification. Input count matrices are normalized independently per cell and log transformed. Existing normalized matrices require --input-scale log1p.
python scripts/download_data.py norman --output data/raw
python -m pivot.cli prepare \
--raw data/raw/NormanWeissman2019_filtered.h5ad \
--dataset norman --split combination --seed 0 --output data/norman_combination
python -m pivot.cli train \
--cache data/norman_combination --config configs/full.json \
--device cuda --output runs/norman_full
python -m pivot.cli evaluate \
--cache data/norman_combination --checkpoint runs/norman_full/best.pt \
--catalog combination --reward cosine --device cuda --output runs/norman_full/test.json
For Replogle, download replogle_k562 and use ReplogleWeissman2022_K562_essential.h5ad with --dataset replogle_k562 --batch-col batch --celltype-col cell_line during preparation. Supported splits are cell, perturbation, gene, and, for Norman, combination. Gene holdouts exclude all training conditions containing the held-out gene. Data documentation defines each partition and the independent response cells used to score nominated actions.
Run the implemented baselines and loss ablations across three seeds:
python scripts/run_matrix.py \
--raw data/raw/NormanWeissman2019_filtered.h5ad \
--dataset norman --split combination --device cuda \
--ablations --output runs/norman_matrix
The experiment matrix maps manuscript comparisons to commands and identifies external integrations that remain to be completed.
Prediction and nomination
A query target is a NumPy array with shape (n_target_cells, n_pca) in the saved cache's PCA basis. The quickstart writes runs/example/target.npy from held-out measured cells. Prediction exports latent coordinates, reconstructed expression, gene order, and source-cell identifiers.
python -m pivot.cli nominate \
--cache runs/example/cache --checkpoint runs/example/model/best.pt \
--target runs/example/target.npy --catalog all --search guidance \
--output runs/example/new_nomination.json
python scripts/plot_results.py runs/example/pivot.json runs/example/ridge.json \
--output runs/example/comparison
Use --search exhaustive to rank the complete catalog, guidance to optimize and project an embedding, or greedy to construct a gene set. Greedy output includes whether the constructed set occurs in the measured catalog. Scores are predictions; measured response evidence is returned by the evaluation workflow.
For a control state $c_0$ and intervention embedding $e$, the endpoint map is $X_\theta(0,1,c_0,e)$. Guidance averages endpoint rewards over source cells and differentiates that average with respect to $e$. The method notes connect the equations to the implementation.
Results, checkpoints, and validation
results/legacy/ and checkpoints/legacy/ preserve the original recorded experiments. Their preprocessing and evaluation conventions differ from the corrected split-first-v1 workflow. The corrected loader checks the cache identity and rejects historical checkpoints. The original source and loading instructions are retained for reproducing the old records.
The corrected example checkpoint is included at checkpoints/example/best.pt, with its matching cache at fixtures/example_cache/. The saved target query is fixtures/example_target.npy. These files support immediate prediction and nomination without retraining.
The corrected small measured-data run exercised preparation, training, reload, evaluation, prediction, exhaustive ranking, guidance, greedy selection, and plot export. Validation details record the executed configurations and remaining full-scale work. No corrected full-scale results are claimed in the blank manuscript tables.
| Directory | Contents |
|---|---|
src/pivot/ |
Data, model, training, inference, and evaluation modules |
configs/ |
Full model, population-loss, and ablation configurations |
scripts/ |
Downloads, complete example, experiment matrix, and figures |
tests/ |
Mathematical and data-isolation regression checks |
fixtures/ |
Small measured dataset and its attribution |
results/ |
Original measurements and executed workflow checks |
checkpoints/ |
Corrected example checkpoint and original trained models |
assets/ |
Shared Figure 1, editable vector source, fonts, and licenses |
docs/ |
Methods, experiments, provenance, and validation |
reproducibility/ |
Original source, upstream revision, bibliography records |