# magic-ai4med/map architecture

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

- Interactive diagram: https://gitdiagram.com/magic-ai4med/map
- Repository: https://github.com/magic-ai4med/map

## Overview

MAP predicts single-cell responses to drugs, including unprofiled drugs, using molecular and cellular representations informed by pharmacological knowledge. The documented workflow prepares Tahoe-100M, OP3 or SciPlex3 data, pre-trains knowledge encoders, trains the response model, then runs inference or evaluation. The graph includes the documented preprocessing and baseline branches; source excerpts give stronger implementation-level evidence for MAP training, evaluation and knowledge-encoder training than for the unsampled branches.

Main entry points: [`MAP/train.py`](https://github.com/magic-ai4med/map/blob/main/MAP/train.py) (MAP training), [`preprocess`](https://github.com/magic-ai4med/map/tree/main/preprocess) (Preprocessing pipeline), [`MAP-KG/train_resume.py`](https://github.com/magic-ai4med/map/blob/main/MAP-KG/train_resume.py) (Encoder pretraining), [`MAP-KG/model/model.py`](https://github.com/magic-ai4med/map/blob/main/MAP-KG/model/model.py) (Drug-gene model)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_prep["Data preparation"]
  node_preprocess["Preprocessing pipeline"]
  node_tahoe_data["Tahoe dataset<br/>[tahoe.py]"]
end

subgraph group_kg["Knowledge encoders"]
  node_kg_data["Knowledge graph data<br/>[dataload.py]"]
  node_kg_model["Drug-gene model<br/>[model.py]"]
  node_kg_train["Encoder pretraining<br/>[train_resume.py]"]
  node_molbart["Molecular encoder<br/>[decoder.py]"]
end

subgraph group_map["Response prediction"]
  node_dataset["Perturbation dataset"]
  node_se_model["Cell-state encoder<br/>[se.py]"]
  node_transformer["Transformer layers"]
  node_map_model["MAP predictor<br/>[model.py]"]
  node_map_train["MAP training<br/>[train.py]"]
  node_demo["Prediction demo<br/>[demo.py]"]
  node_evaluation["Response evaluation<br/>[eval_multi_cell.py]"]
  node_eval_utils["Evaluation metrics<br/>[eval_utils.py]"]
end

subgraph group_baselines["Baseline methods"]
  node_baseline_crisp["CRISP baseline<br/>[train.py]"]
  node_baseline_prnet["PRnet baseline<br/>[PRnet.py]"]
  node_baseline_xpert["XPert baseline<br/>[model_XPert.py]"]
  node_baseline_chemcpa["chemCPA baseline<br/>[model.py]"]
  node_baseline_mean["Mean baseline<br/>[eval_trainmean.py]"]
end

node_raw_data[("Single-cell datasets")]
node_user(("Researcher"))
node_pretrained_weights[("Pretrained weights")]

node_raw_data -->|"prepares via"| node_preprocess
node_preprocess -->|"loads Tahoe"| node_tahoe_data
node_user -->|"starts pretraining"| node_kg_train
node_kg_train -->|"loads batches"| node_kg_data
node_kg_train -->|"trains"| node_kg_model
node_kg_model -->|"uses decoder"| node_molbart
node_user -->|"starts training"| node_map_train
node_preprocess -->|"supplies inputs"| node_dataset
node_map_train -->|"loads batches"| node_dataset
node_map_train -->|"trains"| node_map_model
node_map_model -->|"encodes cell state"| node_se_model
node_se_model -->|"uses layers"| node_transformer
node_map_train -->|"evaluates predictions"| node_evaluation
node_evaluation -->|"uses metrics"| node_eval_utils
node_user -->|"requests prediction"| node_demo
node_pretrained_weights -->|"loads checkpoint"| node_demo
node_user -.->|"runs baseline"| node_baseline_crisp
node_user -.->|"runs baseline"| node_baseline_prnet
node_user -.->|"runs baseline"| node_baseline_xpert
node_user -.->|"runs baseline"| node_baseline_chemcpa
node_user -.->|"runs baseline"| node_baseline_mean

click node_preprocess "https://github.com/magic-ai4med/map/tree/main/preprocess"
click node_tahoe_data "https://github.com/magic-ai4med/map/blob/main/preprocess/data/tahoe.py"
click node_kg_data "https://github.com/magic-ai4med/map/blob/main/MAP-KG/data/dataload.py"
click node_kg_model "https://github.com/magic-ai4med/map/blob/main/MAP-KG/model/model.py"
click node_kg_train "https://github.com/magic-ai4med/map/blob/main/MAP-KG/train_resume.py"
click node_molbart "https://github.com/magic-ai4med/map/blob/main/MAP-KG/model/mega_molbart/decoder.py"
click node_dataset "https://github.com/magic-ai4med/map/blob/main/MAP/data/ds_multi_cell_lora_se.py"
click node_se_model "https://github.com/magic-ai4med/map/blob/main/MAP/model/se.py"
click node_transformer "https://github.com/magic-ai4med/map/blob/main/MAP/model/flash_transformer.py"
click node_map_model "https://github.com/magic-ai4med/map/blob/main/MAP/model/model.py"
click node_map_train "https://github.com/magic-ai4med/map/blob/main/MAP/train.py"
click node_demo "https://github.com/magic-ai4med/map/blob/main/MAP/demo.py"
click node_evaluation "https://github.com/magic-ai4med/map/blob/main/MAP/evaluate/eval_multi_cell.py"
click node_eval_utils "https://github.com/magic-ai4med/map/blob/main/MAP/evaluate/eval_utils.py"
click node_baseline_crisp "https://github.com/magic-ai4med/map/blob/main/MAP/baseline/CRISP/train.py"
click node_baseline_prnet "https://github.com/magic-ai4med/map/blob/main/MAP/baseline/PRnet/models/PRnet.py"
click node_baseline_xpert "https://github.com/magic-ai4med/map/blob/main/MAP/baseline/XPert/models/model_XPert.py"
click node_baseline_chemcpa "https://github.com/magic-ai4med/map/blob/main/MAP/baseline/chemCPA/chemCPA/model.py"
click node_baseline_mean "https://github.com/magic-ai4med/map/blob/main/MAP/baseline/trainmean/eval_trainmean.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_preprocess,node_tahoe_data toneBlue
class node_kg_data,node_kg_model,node_kg_train,node_molbart,node_raw_data,node_pretrained_weights toneAmber
class node_dataset,node_se_model,node_transformer,node_map_model,node_map_train,node_demo,node_evaluation,node_eval_utils toneMint
class node_baseline_crisp,node_baseline_prnet,node_baseline_xpert,node_baseline_chemcpa,node_baseline_mean toneRose
class node_user toneIndigo
```

## Components

### Data preparation

- [Preprocessing pipeline](https://github.com/magic-ai4med/map/tree/main/preprocess): `preprocess`
- [Tahoe dataset](https://github.com/magic-ai4med/map/blob/main/preprocess/data/tahoe.py): `preprocess/data/tahoe.py`

### Knowledge encoders

- [Knowledge graph data](https://github.com/magic-ai4med/map/blob/main/MAP-KG/data/dataload.py): `MAP-KG/data/dataload.py`
- [Drug-gene model](https://github.com/magic-ai4med/map/blob/main/MAP-KG/model/model.py): `MAP-KG/model/model.py`
- [Encoder pretraining](https://github.com/magic-ai4med/map/blob/main/MAP-KG/train_resume.py): `MAP-KG/train_resume.py`
- [Molecular encoder](https://github.com/magic-ai4med/map/blob/main/MAP-KG/model/mega_molbart/decoder.py): `MAP-KG/model/mega_molbart/decoder.py`

### Response prediction

- [Perturbation dataset](https://github.com/magic-ai4med/map/blob/main/MAP/data/ds_multi_cell_lora_se.py): `MAP/data/ds_multi_cell_lora_se.py`
- [Cell-state encoder](https://github.com/magic-ai4med/map/blob/main/MAP/model/se.py): `MAP/model/se.py`
- [Transformer layers](https://github.com/magic-ai4med/map/blob/main/MAP/model/flash_transformer.py): `MAP/model/flash_transformer.py`
- [MAP predictor](https://github.com/magic-ai4med/map/blob/main/MAP/model/model.py): `MAP/model/model.py`
- [MAP training](https://github.com/magic-ai4med/map/blob/main/MAP/train.py): `MAP/train.py`
- [Prediction demo](https://github.com/magic-ai4med/map/blob/main/MAP/demo.py): `MAP/demo.py`
- [Response evaluation](https://github.com/magic-ai4med/map/blob/main/MAP/evaluate/eval_multi_cell.py): `MAP/evaluate/eval_multi_cell.py`
- [Evaluation metrics](https://github.com/magic-ai4med/map/blob/main/MAP/evaluate/eval_utils.py): `MAP/evaluate/eval_utils.py`

### Baseline methods

- [CRISP baseline](https://github.com/magic-ai4med/map/blob/main/MAP/baseline/CRISP/train.py): `MAP/baseline/CRISP/train.py`
- [PRnet baseline](https://github.com/magic-ai4med/map/blob/main/MAP/baseline/PRnet/models/PRnet.py): `MAP/baseline/PRnet/models/PRnet.py`
- [XPert baseline](https://github.com/magic-ai4med/map/blob/main/MAP/baseline/XPert/models/model_XPert.py): `MAP/baseline/XPert/models/model_XPert.py`
- [chemCPA baseline](https://github.com/magic-ai4med/map/blob/main/MAP/baseline/chemCPA/chemCPA/model.py): `MAP/baseline/chemCPA/chemCPA/model.py`
- [Mean baseline](https://github.com/magic-ai4med/map/blob/main/MAP/baseline/trainmean/eval_trainmean.py): `MAP/baseline/trainmean/eval_trainmean.py`

### Other components

- **Single-cell datasets**
- **Researcher**
- **Pretrained weights**

## Connections

- Single-cell datasets → Preprocessing pipeline: prepares via
- Preprocessing pipeline → Tahoe dataset: loads Tahoe (evidence: `preprocess/data/tahoe.py`)
- Researcher → Encoder pretraining: starts pretraining
- Encoder pretraining → Knowledge graph data: loads batches (evidence: `MAP-KG/train_resume.py`)
- Encoder pretraining → Drug-gene model: trains (evidence: `MAP-KG/train_resume.py`)
- Drug-gene model → Molecular encoder: uses decoder
- Researcher → MAP training: starts training
- Preprocessing pipeline → Perturbation dataset: supplies inputs
- MAP training → Perturbation dataset: loads batches (evidence: `MAP/train.py`)
- MAP training → MAP predictor: trains (evidence: `MAP/train.py`)
- MAP predictor → Cell-state encoder: encodes cell state (evidence: `MAP/model/model.py`)
- Cell-state encoder → Transformer layers: uses layers (evidence: `MAP/model/se.py`)
- MAP training → Response evaluation: evaluates predictions (evidence: `MAP/train.py`)
- Response evaluation → Evaluation metrics: uses metrics (evidence: `MAP/evaluate/eval_multi_cell.py`)
- Researcher → Prediction demo: requests prediction
- Pretrained weights → Prediction demo: loads checkpoint
- Researcher → CRISP baseline: runs baseline
- Researcher → PRnet baseline: runs baseline
- Researcher → XPert baseline: runs baseline
- Researcher → chemCPA baseline: runs baseline
- Researcher → Mean baseline: runs baseline