# difra100/pt-rag_iclr architecture

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

- Interactive diagram: https://gitdiagram.com/difra100/pt-rag_iclr
- Repository: https://github.com/difra100/pt-rag_iclr

## Overview

PT-RAG predicts gene-perturbation responses in single-cell data. Its main workflow is to prepare perturbation data, train a state-transition model—optionally with differentiable, cell-type-aware retrieval—and run inference to generate and evaluate predictions. The repository also exposes a distinct embedding workflow for fitting or transforming cell embeddings and querying a LanceDB vector store. The README documents these capabilities; sampled source confirms CLI dispatch, embedding query behavior, data/model components, and prediction-time RAG handling, but does not expose every training or preprocessing call.

Main entry points: [`src/state/_cli/_tx/_predict.py`](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_predict.py) (Prediction Runner), [`src/state/tx/models/state_transition.py`](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/models/state_transition.py) (State Transition Models), [`src/state/_cli/_tx`](https://github.com/difra100/pt-rag_iclr/tree/main/src/state/_cli/_tx) (TX Commands), [`src/state/_cli/_emb`](https://github.com/difra100/pt-rag_iclr/tree/main/src/state/_cli/_emb) (Embedding Commands)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_cli["Command Interface"]
  node_cli["State CLI<br/>[__main__.py]"]
  node_txcommands["TX Commands"]
  node_embcommands["Embedding Commands"]
  node_training["Training Runner<br/>[_train.py]"]
  node_inference["Prediction Runner<br/>[_predict.py]"]
  node_preprocess["Data Preparation"]
end

subgraph group_tx["Perturbation Modeling"]
  node_pertdata["Perturbation Dataset"]
  node_transition["State Transition Models"]
  node_modelzoo["Model Variants"]
  node_scgpt["scGPT Components"]
  node_rag["Differentiable Retrieval"]
  node_predictions["Predicted Responses"]
end

subgraph group_emb["Embedding Workflow"]
  node_embfit["Embedding Training<br/>[trainer.py]"]
  node_embmodel["Embedding Model<br/>[model.py]"]
  node_embdata["Embedding Data Loader<br/>[loader.py]"]
  node_embinference["Embedding Inference<br/>[inference.py]"]
  node_query["Vector Query<br/>[_query.py]"]
end

subgraph group_integrations["Evaluation And Storage"]
  node_vectordb[("State Vector DB<br/>[vectordb.py]")]
  node_lancedb[("LanceDB")]
  node_gge["GenGeneEval"]
  node_metrics["Evaluation Metrics<br/>[_predict.py]"]
end

node_user(("Researcher"))

node_user -->|"invokes"| node_cli
node_cli -->|"dispatches"| node_txcommands
node_cli -->|"dispatches"| node_embcommands
node_txcommands -->|"selects train"| node_training
node_txcommands -->|"selects predict"| node_inference
node_txcommands -->|"selects preparation"| node_preprocess
node_preprocess -->|"prepares inputs"| node_pertdata
node_training -->|"loads batches"| node_pertdata
node_training -->|"trains"| node_transition
node_transition -->|"learns retrieval"| node_rag
node_transition -->|"uses configured model"| node_modelzoo
node_transition -->|"uses components"| node_scgpt
node_inference -->|"sets up test data"| node_pertdata
node_inference -->|"loads checkpoint"| node_transition
node_inference -->|"writes responses"| node_predictions
node_inference -->|"evaluates with"| node_gge
node_gge -->|"returns metrics"| node_metrics
node_embcommands -->|"selects fit"| node_embfit
node_embfit -->|"loads data"| node_embdata
node_embfit -->|"trains"| node_embmodel
node_embcommands -->|"selects transform"| node_embinference
node_embinference -->|"loads cells"| node_embdata
node_embinference -->|"embeds cells"| node_embmodel
node_embinference -.->|"updates index"| node_vectordb
node_vectordb -.->|"persists vectors"| node_lancedb
node_embcommands -->|"selects query"| node_query
node_query -->|"searches vectors"| node_vectordb
node_query -->|"queries table"| node_lancedb

click node_cli "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/__main__.py"
click node_txcommands "https://github.com/difra100/pt-rag_iclr/tree/main/src/state/_cli/_tx"
click node_embcommands "https://github.com/difra100/pt-rag_iclr/tree/main/src/state/_cli/_emb"
click node_training "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_train.py"
click node_inference "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_predict.py"
click node_preprocess "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_preprocess_train.py"
click node_pertdata "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/data/dataset/scgpt_perturbation_dataset.py"
click node_transition "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/models/state_transition.py"
click node_modelzoo "https://github.com/difra100/pt-rag_iclr/tree/main/src/state/tx/models"
click node_scgpt "https://github.com/difra100/pt-rag_iclr/tree/main/src/state/tx/models/scgpt"
click node_rag "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/models/state_transition.py"
click node_embfit "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/train/trainer.py"
click node_embmodel "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/nn/model.py"
click node_embdata "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/data/loader.py"
click node_embinference "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/inference.py"
click node_query "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_emb/_query.py"
click node_vectordb "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/vectordb.py"
click node_metrics "https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_predict.py"

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_cli,node_txcommands,node_embcommands,node_training,node_inference,node_preprocess toneBlue
class node_pertdata,node_transition,node_modelzoo,node_scgpt,node_rag,node_predictions toneAmber
class node_embfit,node_embmodel,node_embdata,node_embinference,node_query toneMint
class node_vectordb,node_lancedb,node_gge,node_metrics toneRose
class node_user toneIndigo
```

## Components

### Command Interface

- [State CLI](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/__main__.py): `src/state/__main__.py`
- [TX Commands](https://github.com/difra100/pt-rag_iclr/tree/main/src/state/_cli/_tx): `src/state/_cli/_tx`
- [Embedding Commands](https://github.com/difra100/pt-rag_iclr/tree/main/src/state/_cli/_emb): `src/state/_cli/_emb`
- [Training Runner](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_train.py): `src/state/_cli/_tx/_train.py`
- [Prediction Runner](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_predict.py): `src/state/_cli/_tx/_predict.py`
- [Data Preparation](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_preprocess_train.py): `src/state/_cli/_tx/_preprocess_train.py`

### Perturbation Modeling

- [Perturbation Dataset](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/data/dataset/scgpt_perturbation_dataset.py): `src/state/tx/data/dataset/scgpt_perturbation_dataset.py`
- [State Transition Models](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/models/state_transition.py): `src/state/tx/models/state_transition.py`
- [Model Variants](https://github.com/difra100/pt-rag_iclr/tree/main/src/state/tx/models): `src/state/tx/models`
- [scGPT Components](https://github.com/difra100/pt-rag_iclr/tree/main/src/state/tx/models/scgpt): `src/state/tx/models/scgpt`
- [Differentiable Retrieval](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/tx/models/state_transition.py): `src/state/tx/models/state_transition.py`
- **Predicted Responses**

### Embedding Workflow

- [Embedding Training](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/train/trainer.py): `src/state/emb/train/trainer.py`
- [Embedding Model](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/nn/model.py): `src/state/emb/nn/model.py`
- [Embedding Data Loader](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/data/loader.py): `src/state/emb/data/loader.py`
- [Embedding Inference](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/inference.py): `src/state/emb/inference.py`
- [Vector Query](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_emb/_query.py): `src/state/_cli/_emb/_query.py`

### Evaluation And Storage

- [State Vector DB](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/emb/vectordb.py): `src/state/emb/vectordb.py`
- **LanceDB**
- **GenGeneEval**
- [Evaluation Metrics](https://github.com/difra100/pt-rag_iclr/blob/main/src/state/_cli/_tx/_predict.py): `src/state/_cli/_tx/_predict.py`

### Other components

- **Researcher**

## Connections

- Researcher → State CLI: invokes
- State CLI → TX Commands: dispatches
- State CLI → Embedding Commands: dispatches
- TX Commands → Training Runner: selects train
- TX Commands → Prediction Runner: selects predict
- TX Commands → Data Preparation: selects preparation
- Data Preparation → Perturbation Dataset: prepares inputs
- Training Runner → Perturbation Dataset: loads batches
- Training Runner → State Transition Models: trains
- State Transition Models → Differentiable Retrieval: learns retrieval
- State Transition Models → Model Variants: uses configured model
- State Transition Models → scGPT Components: uses components
- Prediction Runner → Perturbation Dataset: sets up test data
- Prediction Runner → State Transition Models: loads checkpoint
- Prediction Runner → Predicted Responses: writes responses
- Prediction Runner → GenGeneEval: evaluates with
- GenGeneEval → Evaluation Metrics: returns metrics
- Embedding Commands → Embedding Training: selects fit
- Embedding Training → Embedding Data Loader: loads data
- Embedding Training → Embedding Model: trains
- Embedding Commands → Embedding Inference: selects transform
- Embedding Inference → Embedding Data Loader: loads cells
- Embedding Inference → Embedding Model: embeds cells
- Embedding Inference → State Vector DB: updates index
- State Vector DB → LanceDB: persists vectors
- Embedding Commands → Vector Query: selects query
- Vector Query → State Vector DB: searches vectors
- Vector Query → LanceDB: queries table