How saigum/stateexperiments works
Purpose and entry points
State is a Python ML package for single-cell perturbation response prediction (“state transitions”) and cell embedding pretraining/inference. The installed state command dispatches two product surfaces: tx and emb, rooted at src/state/__main__.py, src/state/_cli/_tx/, and src/state/_cli/_emb/.
Transition-model workflow
The tx CLI exposes preprocessing, training, prediction/evaluation, and inference stages: src/state/_cli/_tx/_preprocess_train.py, src/state/_cli/_tx/_train.py, and src/state/_cli/_tx/_infer.py. Training consumes AnnData/H5AD datasets plus TOML-defined dataset and zeroshot/fewshot splits; inference applies a checkpoint to prepared control-state inputs. Outputs include trained checkpoints and predicted AnnData-style results.
Transition data boundary
Training preprocessing normalizes/log-transforms expression and produces HVG embeddings; inference preprocessing replaces perturbed expression with controls while retaining perturbation labels. These distinct boundaries are implemented in src/state/_cli/_tx/_preprocess_train.py and src/state/_cli/_tx/_preprocess_infer.py. Dataset adaptation for scGPT perturbation models resides in src/state/tx/data/dataset/scgpt_perturbation_dataset.py. AnnData/Scanpy-compatible H5AD is the primary interchange format.
Transition model family
tx supports interchangeable prediction architectures selected through model configuration: the native State transition model, set/embedding-based models, mean baselines, decoder-only models, CPA, scGPT, and scVI. Core implementations are under src/state/tx/models/state_transition.py, src/state/tx/models/cpa/_model.py, and src/state/tx/models/scgpt/lightning_model.py; model selection defaults live in src/state/configs/model/. This is a model-pluggable training system rather than one fixed network.
Embedding workflow
The emb branch trains transformer-based cell embedding models, transforms H5AD inputs using checkpoints, evaluates embeddings, and can fine-tune a decoder. Its CLI stages are src/state/_cli/_emb/_fit.py, src/state/_cli/_emb/_transform.py, and src/state/_cli/_emb/_query.py. The embedding model/data/training boundary is represented by src/state/emb/nn/model.py, src/state/emb/data/loader.py, and src/state/emb/train/trainer.py.
Retrieval and external integrations
Embedding transforms may append vectors to LanceDB and query nearest cells through src/state/emb/vectordb.py; this is an optional vector-database dependency. Experiment configuration includes Weights & Biases defaults in src/state/configs/wandb/default.yaml, making W&B a training-observability integration. Checkpoints may be sourced from Hugging Face model folders, while associated external projects provide evaluation (cell-eval) and broader data loading/preprocessing (cell-load).
Runtime and deployment
The project is a typed Python package (src/state/py.typed) distributed through PyPI/uv, with Python version pinned by .python-version and dependencies/package metadata in pyproject.toml. GPU-oriented execution is supported by the Singularity recipe at singularity.def (including --nv usage). Release automation is defined in .github/workflows/release.yml.
Main entry points: _train.py Transition training · _infer.py Transition inference · state_transition.py Transition model family · trainer.py Embedding trainer
Components 25
Command surfaces
- state command · Dispatches tx and emb product surfaces.
src/state/__main__.py - tx CLI
- emb CLI
Transition prediction
- AnnData / H5AD · Primary Scanpy-compatible dataset boundary.
- Training preprocessing · Normalizes and log-transforms expression and produces HVG embeddings.
src/state/_cli/_tx/_preprocess_train.py - Inference preprocessing · Replaces perturbed expression with controls while retaining perturbation labels.
src/state/_cli/_tx/_preprocess_infer.py - Transition training · Consumes prepared data plus TOML dataset and split definitions.
src/state/_cli/_tx/_train.py - Transition inference
src/state/_cli/_tx/_infer.py - scGPT dataset adapter
src/state/tx/data/dataset/scgpt_perturbation_dataset.py - Model selection config · Selects among interchangeable transition architectures.
- Transition model family · Native State, set/embedding, baselines, decoder-only, CPA, scGPT, and scVI models.
src/state/tx/models/state_transition.py - CPA model
src/state/tx/models/cpa/_model.py - scGPT Lightning model
src/state/tx/models/scgpt/lightning_model.py - Transition checkpoints
Cell embeddings
- Embedding fit
src/state/_cli/_emb/_fit.py - Embedding transform
src/state/_cli/_emb/_transform.py - Embedding query
src/state/_cli/_emb/_query.py - Embedding data loader
src/state/emb/data/loader.py - Embedding transformer
src/state/emb/nn/model.py - Embedding trainer
src/state/emb/train/trainer.py - LanceDB · Stores vectors and serves nearest-cell retrieval.
src/state/emb/vectordb.py
Distribution & integrations
- Weights & Biases
src/state/configs/wandb/default.yaml - Python package · Typed package with pinned Python and declared dependencies.
pyproject.toml - Singularity runtime · Supports GPU execution with --nv.
singularity.def - Release automation
.github/workflows/release.yml
Connections 31
- state command → tx CLI: dispatches
- state command → emb CLI: dispatches
- tx CLI → Training preprocessing: preprocess
- tx CLI → Transition training: train
- tx CLI → Inference preprocessing: prepare inference
- tx CLI → Transition inference: infer
- AnnData / H5AD → Training preprocessing: training input
- AnnData / H5AD → Inference preprocessing: control-state input
- Training preprocessing → Transition training: prepared training data
- Inference preprocessing → Transition inference: prepared controls
- Model selection config → Transition training: selects model
- Transition training → Transition model family: trains
- Transition model family → CPA model: includes
- Transition model family → scGPT Lightning model: includes
- scGPT dataset adapter → scGPT Lightning model: adapts data for
- Transition training → Transition checkpoints: writes
- Transition checkpoints → Transition inference: loads
- Transition inference → AnnData / H5AD: predicted results
- emb CLI → Embedding fit: fit
- emb CLI → Embedding transform: transform
- emb CLI → Embedding query: query / evaluate
- AnnData / H5AD → Embedding data loader: embedding input
- Embedding data loader → Embedding trainer: batches
- Embedding trainer → Embedding transformer: trains
- Embedding transformer → Embedding transform: checkpoint inference
- Embedding transform → LanceDB: optionally appends vectors
- Embedding query → LanceDB: nearest-cell search
- Transition training → Weights & Biases: logs experiments
- Embedding trainer → Weights & Biases: logs experiments
- Python package → Singularity runtime: deployed in
- Release automation → Python package: publishes