# saigum/stateexperiments architecture

> Architecture diagram of the GitHub repository saigum/stateexperiments, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-11). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/saigum/stateexperiments
- Repository: https://github.com/saigum/stateexperiments

## Overview

#### 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: [`src/state/_cli/_tx/_train.py`](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_train.py) (Transition training), [`src/state/_cli/_tx/_infer.py`](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_infer.py) (Transition inference), [`src/state/tx/models/state_transition.py`](https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/state_transition.py) (Transition model family), [`src/state/emb/train/trainer.py`](https://github.com/saigum/stateexperiments/blob/main/src/state/emb/train/trainer.py) (Embedding trainer)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_cli["Command surfaces"]
  node_state_cli{{"state command<br/>Python CLI dispatcher<br/>[__main__.py]"}}
  node_tx_cli["tx CLI<br/>transition commands"]
  node_emb_cli["emb CLI<br/>embedding commands"]
end

subgraph group_tx["Transition prediction"]
  node_h5ad[("AnnData / H5AD<br/>data interchange")]
  node_tx_train_preprocess["Training preprocessing<br/>expression preparation"]
  node_tx_infer_preprocess["Inference preprocessing<br/>control-state preparation"]
  node_tx_train["Transition training<br/>trainer stage<br/>[_train.py]"]
  node_tx_infer["Transition inference<br/>prediction &amp; evaluation<br/>[_infer.py]"]
  node_scgpt_adapter["scGPT dataset adapter<br/>dataset adaptation"]
  node_model_configs["Model selection config<br/>model defaults"]
  node_tx_models["Transition model family<br/>pluggable predictors"]
  node_cpa_model["CPA model<br/>transition architecture<br/>[_model.py]"]
  node_scgpt_model["scGPT Lightning model<br/>transition architecture<br/>[lightning_model.py]"]
  node_tx_checkpoint["Transition checkpoints<br/>trained artifacts"]
end

subgraph group_emb["Cell embeddings"]
  node_emb_fit["Embedding fit<br/>training command<br/>[_fit.py]"]
  node_emb_transform["Embedding transform<br/>inference command<br/>[_transform.py]"]
  node_emb_query["Embedding query<br/>evaluation &amp; retrieval command<br/>[_query.py]"]
  node_emb_data["Embedding data loader<br/>data boundary<br/>[loader.py]"]
  node_emb_model["Embedding transformer<br/>cell embedding model<br/>[model.py]"]
  node_emb_trainer["Embedding trainer<br/>training runtime<br/>[trainer.py]"]
  node_vector_db[("LanceDB<br/>optional vector database<br/>[vectordb.py]")]
end

subgraph group_ops["Distribution &amp; integrations"]
  node_wandb["Weights &amp; Biases<br/>experiment observability<br/>[default.yaml]"]
  node_package["Python package<br/>PyPI / uv distribution<br/>[pyproject.toml]"]
  node_gpu_image["Singularity runtime<br/>GPU deployment<br/>[singularity.def]"]
  node_release["Release automation<br/>CI release workflow<br/>[release.yml]"]
end

node_state_cli -->|"dispatches"| node_tx_cli
node_state_cli -->|"dispatches"| node_emb_cli
node_tx_cli -->|"preprocess"| node_tx_train_preprocess
node_tx_cli -->|"train"| node_tx_train
node_tx_cli -->|"prepare inference"| node_tx_infer_preprocess
node_tx_cli -->|"infer"| node_tx_infer
node_h5ad -->|"training input"| node_tx_train_preprocess
node_h5ad -->|"control-state input"| node_tx_infer_preprocess
node_tx_train_preprocess -->|"prepared training data"| node_tx_train
node_tx_infer_preprocess -->|"prepared controls"| node_tx_infer
node_model_configs -->|"selects model"| node_tx_train
node_tx_train -->|"trains"| node_tx_models
node_tx_models -.->|"includes"| node_cpa_model
node_tx_models -.->|"includes"| node_scgpt_model
node_scgpt_adapter -->|"adapts data for"| node_scgpt_model
node_tx_train -->|"writes"| node_tx_checkpoint
node_tx_checkpoint -->|"loads"| node_tx_infer
node_tx_infer -->|"predicted results"| node_h5ad
node_emb_cli -->|"fit"| node_emb_fit
node_emb_cli -->|"transform"| node_emb_transform
node_emb_cli -->|"query / evaluate"| node_emb_query
node_h5ad -->|"embedding input"| node_emb_data
node_emb_data -->|"batches"| node_emb_trainer
node_emb_trainer -->|"trains"| node_emb_model
node_emb_model -->|"checkpoint inference"| node_emb_transform
node_emb_transform -.->|"optionally appends vectors"| node_vector_db
node_emb_query -.->|"nearest-cell search"| node_vector_db
node_tx_train -.->|"logs experiments"| node_wandb
node_emb_trainer -.->|"logs experiments"| node_wandb
node_package -.->|"deployed in"| node_gpu_image
node_release -->|"publishes"| node_package

click node_state_cli "https://github.com/saigum/stateexperiments/blob/main/src/state/__main__.py"
click node_tx_train_preprocess "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_preprocess_train.py"
click node_tx_infer_preprocess "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_preprocess_infer.py"
click node_tx_train "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_train.py"
click node_tx_infer "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_infer.py"
click node_scgpt_adapter "https://github.com/saigum/stateexperiments/blob/main/src/state/tx/data/dataset/scgpt_perturbation_dataset.py"
click node_tx_models "https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/state_transition.py"
click node_cpa_model "https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/cpa/_model.py"
click node_scgpt_model "https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/scgpt/lightning_model.py"
click node_emb_fit "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_emb/_fit.py"
click node_emb_transform "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_emb/_transform.py"
click node_emb_query "https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_emb/_query.py"
click node_emb_data "https://github.com/saigum/stateexperiments/blob/main/src/state/emb/data/loader.py"
click node_emb_model "https://github.com/saigum/stateexperiments/blob/main/src/state/emb/nn/model.py"
click node_emb_trainer "https://github.com/saigum/stateexperiments/blob/main/src/state/emb/train/trainer.py"
click node_vector_db "https://github.com/saigum/stateexperiments/blob/main/src/state/emb/vectordb.py"
click node_wandb "https://github.com/saigum/stateexperiments/blob/main/src/state/configs/wandb/default.yaml"
click node_package "https://github.com/saigum/stateexperiments/blob/main/pyproject.toml"
click node_gpu_image "https://github.com/saigum/stateexperiments/blob/main/singularity.def"
click node_release "https://github.com/saigum/stateexperiments/blob/main/.github/workflows/release.yml"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_state_cli,node_tx_cli,node_emb_cli toneBlue
class node_h5ad,node_tx_train_preprocess,node_tx_infer_preprocess,node_tx_train,node_tx_infer,node_scgpt_adapter,node_model_configs,node_tx_models,node_cpa_model,node_scgpt_model,node_tx_checkpoint toneAmber
class node_emb_fit,node_emb_transform,node_emb_query,node_emb_data,node_emb_model,node_emb_trainer,node_vector_db toneMint
class node_wandb,node_package,node_gpu_image,node_release toneRose
```

## Components

### Command surfaces

- [state command](https://github.com/saigum/stateexperiments/blob/main/src/state/__main__.py): Python CLI dispatcher · Dispatches tx and emb product surfaces. · `src/state/__main__.py`
- **tx CLI**: transition commands
- **emb CLI**: embedding commands

### Transition prediction

- **AnnData / H5AD**: data interchange · Primary Scanpy-compatible dataset boundary.
- [Training preprocessing](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_preprocess_train.py): expression preparation · Normalizes and log-transforms expression and produces HVG embeddings. · `src/state/_cli/_tx/_preprocess_train.py`
- [Inference preprocessing](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_preprocess_infer.py): control-state preparation · Replaces perturbed expression with controls while retaining perturbation labels. · `src/state/_cli/_tx/_preprocess_infer.py`
- [Transition training](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_train.py): trainer stage · Consumes prepared data plus TOML dataset and split definitions. · `src/state/_cli/_tx/_train.py`
- [Transition inference](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_tx/_infer.py): prediction & evaluation · `src/state/_cli/_tx/_infer.py`
- [scGPT dataset adapter](https://github.com/saigum/stateexperiments/blob/main/src/state/tx/data/dataset/scgpt_perturbation_dataset.py): dataset adaptation · `src/state/tx/data/dataset/scgpt_perturbation_dataset.py`
- **Model selection config**: model defaults · Selects among interchangeable transition architectures.
- [Transition model family](https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/state_transition.py): pluggable predictors · Native State, set/embedding, baselines, decoder-only, CPA, scGPT, and scVI models. · `src/state/tx/models/state_transition.py`
- [CPA model](https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/cpa/_model.py): transition architecture · `src/state/tx/models/cpa/_model.py`
- [scGPT Lightning model](https://github.com/saigum/stateexperiments/blob/main/src/state/tx/models/scgpt/lightning_model.py): transition architecture · `src/state/tx/models/scgpt/lightning_model.py`
- **Transition checkpoints**: trained artifacts

### Cell embeddings

- [Embedding fit](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_emb/_fit.py): training command · `src/state/_cli/_emb/_fit.py`
- [Embedding transform](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_emb/_transform.py): inference command · `src/state/_cli/_emb/_transform.py`
- [Embedding query](https://github.com/saigum/stateexperiments/blob/main/src/state/_cli/_emb/_query.py): evaluation & retrieval command · `src/state/_cli/_emb/_query.py`
- [Embedding data loader](https://github.com/saigum/stateexperiments/blob/main/src/state/emb/data/loader.py): data boundary · `src/state/emb/data/loader.py`
- [Embedding transformer](https://github.com/saigum/stateexperiments/blob/main/src/state/emb/nn/model.py): cell embedding model · `src/state/emb/nn/model.py`
- [Embedding trainer](https://github.com/saigum/stateexperiments/blob/main/src/state/emb/train/trainer.py): training runtime · `src/state/emb/train/trainer.py`
- [LanceDB](https://github.com/saigum/stateexperiments/blob/main/src/state/emb/vectordb.py): optional vector database · Stores vectors and serves nearest-cell retrieval. · `src/state/emb/vectordb.py`

### Distribution & integrations

- [Weights & Biases](https://github.com/saigum/stateexperiments/blob/main/src/state/configs/wandb/default.yaml): experiment observability · `src/state/configs/wandb/default.yaml`
- [Python package](https://github.com/saigum/stateexperiments/blob/main/pyproject.toml): PyPI / uv distribution · Typed package with pinned Python and declared dependencies. · `pyproject.toml`
- [Singularity runtime](https://github.com/saigum/stateexperiments/blob/main/singularity.def): GPU deployment · Supports GPU execution with --nv. · `singularity.def`
- [Release automation](https://github.com/saigum/stateexperiments/blob/main/.github/workflows/release.yml): CI release workflow · `.github/workflows/release.yml`

## Connections

- 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