---
language: en
license: mit
tags:
  - biology
  - single-cell
  - perturbation-prediction
  - flow-matching
  - pytorch
---

# PIVOT

![PIVOT overview](assets/figure1.png)

**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](docs/results.md) 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.

```bash
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](https://zenodo.org/records/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`.

```bash
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](docs/data.md) defines each partition and the independent response cells used to score nominated actions.

Run the implemented baselines and loss ablations across three seeds:

```bash
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](docs/experiments.md) 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.

```bash
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](docs/method.md) 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](docs/validation.md) 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 |