State and the shift from single-cell atlases to perturbation prediction

State targets the response to a perturbation, not the state of a cell at rest

The repository description calls State a machine learning model that predicts cellular perturbation response across diverse contexts. That word, perturbation, is the whole scope. The project sits downstream of single-cell atlases: the input is a matrix of counts where cells have been exposed to a known genetic or chemical intervention, and the output is a prediction of what the transcriptome looks like after a different or unseen intervention.

The audience is therefore computational biologists with a specific asset, an h5ad file from a perturbation experiment, and a specific question, can this model extrapolate to a cell type or a target gene it has not seen. The README points to a paper, a Virtual Cell Challenge Colab, and several Colab notebooks for training on the Replogle-Nadig dataset and running inference on Tahoe-100M, so the intended entry points are hosted notebooks rather than a local build.

The distinction matters for expectations. State is not a general embedding model for cell similarity. It predicts a cellular state change, which is a harder claim and a narrower one.

One `state` command, two model families

The CLI is organised around two subcommand families. `state tx` drives the State Transition model, and `state emb` drives the State Embedding model. The README describes the purpose of the first as training and running perturbation prediction models, and the second as pretraining cell embeddings or embedding new datasets.

Both are reached through the same entry point. `pyproject.toml` maps the script to `state.__main__:main`, and the README explains that an installation made with `uv tool install` is invoked as `state ...`, while a checkout is driven with `uv run state ...`. Subcommand help is available through `state --help`, `state tx --help` and `state emb --help`.

This split is the practical shape of the project. A prediction workflow goes through the tx path with preprocessing, training, evaluation and inference. An embedding workflow goes through the emb path with its own preprocessing step, a fit stage, and a transform stage that the README does not finish describing, since the section is cut off mid sentence.

A TOML file decides what training, zero-shot and few-shot mean

Before any model runs, the data splits have to be declared. The file referenced by `data.kwargs.toml_config_path` defines dataset paths and splits, with a required `[datasets]` section mapping dataset names to directories of h5ad files and a required `[training]` section selecting which datasets participate in training.

Two optional sections control the evaluation design. `[zeroshot]` holds out entire cell types for validation and test, which measures generalisation to a context the model has never trained on. `[fewshot]` holds out specific perturbations within specific cell types, which measures generalisation to a new target gene in a familiar context. The README is explicit that the key format is `"dataset.cell_type"` in both sections, and that anything not listed defaults to training.

A minimal configuration looks like this:

[datasets]
replogle = "/data/replogle/"

[training]
replogle = "train"

[zeroshot]
"replogle.jurkat" = "test"

[fewshot]
[fewshot."replogle.k562"]
val = ["AARS"]
test = ["NUP107"]

That file is the real interface of the project. A dataset you forget to list silently becomes training data, so the split file deserves the same attention as any hyperparameter.

preprocess_train fixes the feature space at highly variable genes

The first step of the transition workflow prepares training h5ad files by normalising counts, applying log1p and selecting highly variable genes. The resulting matrix is stored in `.obsm["X_hvg"]` and written to the path given by `--output`.

The gene count is a flag, not a constant. In the example the README gives, `--num_hvgs 2000` sets it:

state tx preprocess_train \
  --adata /path/to/raw_data.h5ad \
  --output /path/to/preprocessed_training_data.h5ad \
  --num_hvgs 2000

Two other knobs shape what the model sees. `data.kwargs.downsample` downsamples counts during loading, where a value of one or less keeps that fraction by binomial sampling and a value above one targets that read depth per cell for `output_space=all`. `data.kwargs.downsample_cells` caps the number of cells loaded per `(cell_type, perturbation[, batch])` group while leaving smaller groups unchanged, which is the lever for balancing an experiment where one perturbation was measured far more deeply than the others.

The README does not document the normalisation target, the minimum count filter or how ties are broken when selecting highly variable genes, so preprocessing choices are yours to verify against the output matrix.

Training is a Hydra command line, so every run is a config diff

The `train` subcommand takes Hydra overrides rather than a config file argument. Dataset selection, feature key, perturbation column, cell type key, step count, batch size, model choice, output directory and run name are all set on the command line:

state tx train \
  data.kwargs.toml_config_path=examples/fewshot.toml \
  data.kwargs.embed_key=X_hvg \
  data.kwargs.pert_col=target_gene \
  data.kwargs.cell_type_key=cell_type \
  training.max_steps=40000 \
  training.batch_size=8 \
  model=state \
  output_dir="$HOME/state" \
  name="test"

`output_dir` and `name` together form the run directory, and the README notes that the prediction step expects to find `config.yaml` and a `checkpoints/` directory inside it. That layout is what makes a run self describing: the saved config and the checkpoint stay together, and evaluation reads the same file.

Evaluation then points at the run directory with `state tx predict --output-dir $HOME/state/test --checkpoint final.ckpt`, using `cell-eval` metrics, or simply running prediction with `--predict-only`. A `--toml` flag evaluates on a different dataset or split configuration than the one saved with the run, which is the mechanism for scoring a trained model against a split you did not train on.

Inference on new data needs an explicit feature key and perturbation column

The `infer` subcommand runs the model on data that was not necessarily part of the training TOML. It needs the run directory, a checkpoint, and an AnnData file:

state tx infer \
  --model-dir /path/to/run \
  --checkpoint /path/to/run/checkpoints/final.ckpt \
  --adata /path/to/preprocessed_data.h5ad \
  --pert-col gene \
  --embed-key X_hvg \
  --output /path/to/output.h5ad

Two flags carry the assumptions from training. `--embed-key X_hvg` tells the model which matrix in `.obsm` to read, and it is required when the model was trained on highly variable gene features. `--pert-col` names the column holding the perturbation label. If the output path ends with `.npy`, only the predictions matrix is written and no h5ad file is produced, which is the convenient path when you want the raw matrix and nothing else.

The failure mode here is quiet. A mismatch between the feature key used at training and the one given here, or a perturbation column that does not exist in your file, produces a mismatch that no part of the README describes. Checking the run's `config.yaml` against your inference flags is the only cross-check the documentation offers.

The embedding path builds a profile before it fits anything

`state emb preprocess` scans your train and validation datasets and writes a profile that later stages reference. It expects CSV files with `species,path,names` columns, where each `names` value becomes a dataset identifier in the config:

species,path,names
human,/path/ds1.h5ad,ds1
human,/path/ds2.h5ad,ds2

The step auto-detects the gene name field in `var`, checking `_index`, `gene_name`, `gene_symbols`, `feature_name`, `gene_id` and `symbol` in that order. It uses `--all-embeddings` if you provide a torch dictionary of `{gene_name: embedding}`, otherwise it builds one-hot embeddings. Outputs are `all_embeddings_{profile}.pt`, `ds_emb_mapping_{profile}.torch`, `valid_genes_masks_{profile}.torch`, plus `train_{profile}.csv` and `val_{profile}.csv` written into `--output-dir`.

By default it edits `src/state/configs/state-defaults.yaml` in place, which is why the README recommends copying that file first and passing `--config-file`, since preprocess only updates the embeddings and dataset sections and expects a complete config. The `fit` stage then needs `embeddings.current` and `dataset.current` pointed at the profile, and the README recommends the same `--config-file` route.

Separate repositories for data loading and evaluation, and a licence that needs reading

State does not stand alone. The README lists `cell-load` for dataloaders and preprocessing, and `cell-eval` for the model evaluation framework, both in separate Arc Institute repositories. The dependency list in `pyproject.toml` confirms the coupling: `cell-load>=0.10.4` and `cell-eval>=0.8.1` sit alongside `anndata`, `scanpy`, `torch`, `hydra-core`, `wandb`, `geomloss`, `transformers` and `peft`. An optional `vectordb` extra pulls in `lancedb`. Python is pinned to a narrow window, `>=3.11,<3.13`, and the build backend is hatchling.

Licensing is the item to look at first. The repository is published under a custom licence that GitHub cannot classify, and it ships two extra documents at the root, `MODEL_LICENSE.md` and `MODEL_ACCEPTABLE_USE_POLICY.md`, which suggests the weights and the code are governed separately. Read both before use in a commercial or institutional setting.

On maintenance, the last push was on 2026-07-24, and the most recent tagged releases are v0.11.0 and v0.11.1, both dated 2026-06-29, with v0.10.2 from 2026-01-23. The version in `pyproject.toml` reads 0.11.3, so the published package version has moved past the tags visible on the repository. The repository also ships `singularity.def`, `state_preprint_toml_files/` and a `tests/` directory, and the example TOMLs for zero-shot, few-shot and mixed setups are in `examples/`.

Editorial conclusion

State is worth the setup effort for groups that hold perturbation datasets with defined cell types and a held-out evaluation design, and that can describe their splits in a TOML file before training. It is the wrong choice for a quick single dataset experiment, for anyone without AnnData h5ad files and GPU time, or for a purely exploratory analysis with no perturbation labels. Verify first that your gene identifiers match the column names the preprocessing step expects, because the embedding path auto-detects the gene name field and writes masks for genes it cannot resolve, and confirm the terms in MODEL_LICENSE.md and MODEL_ACCEPTABLE_USE_POLICY.md before using the released weights.

Frequently asked questions

Official sources

1. ArcInstitute/state on GitHub
2. Issues
3. README
4. Releases