# nyumedml/gnn_for_ehr architecture

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

- Interactive diagram: https://gitdiagram.com/nyumedml/gnn_for_ehr
- Repository: https://github.com/nyumedml/gnn_for_ehr

## Overview

This repository implements graph-based EHR representation learning for predicting disease outcomes from ICU records. Researchers preprocess MIMIC or eICU encounter data into labeled feature matrices, then train and evaluate a variational graph neural network and save model checkpoints. The README also describes interpretation through singular-value analysis and representation clustering; that capability is documented, but its implementation is not identifiable from the sampled source, so no runtime wiring is asserted for it.

Main entry points: [`train.py`](https://github.com/nyumedml/gnn_for_ehr/blob/master/train.py) (Training and evaluation loop), [`data/mimic/preprocess_mimic.py`](https://github.com/nyumedml/gnn_for_ehr/blob/master/data/mimic/preprocess_mimic.py) (MIMIC preprocessing), [`data/eicu/preprocess_eicu.py`](https://github.com/nyumedml/gnn_for_ehr/blob/master/data/eicu/preprocess_eicu.py) (eICU preprocessing), [`data`](https://github.com/nyumedml/gnn_for_ehr/tree/master/data) (Processed splits)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_preparation["EHR preparation"]
  node_mimic["MIMIC preprocessing"]
  node_eicu["eICU preprocessing<br/>[preprocess_eicu.py]"]
end

subgraph group_learning["Model learning"]
  node_trainer["Training and evaluation loop<br/>[train.py]"]
  node_batching["Batch loading<br/>[utils.py]"]
  node_model["Variational graph model<br/>[model.py]"]
  node_interpretation["Singular-value analysis"]
end

subgraph group_assets["Data and results"]
  node_processed["Processed splits"]
  node_metrics["Validation AUPRC<br/>[utils.py]"]
  node_checkpoints["Model checkpoints"]
end

node_researcher(("Researcher"))
node_rawehr["Raw EHR records"]
node_prediction["Disease-outcome predictions"]

node_researcher -->|"runs preprocessing"| node_mimic
node_researcher -->|"runs preprocessing"| node_eicu
node_rawehr -->|"reads records"| node_mimic
node_rawehr -->|"reads records"| node_eicu
node_mimic -->|"writes splits"| node_processed
node_eicu -->|"writes splits"| node_processed
node_researcher -->|"starts training"| node_trainer
node_processed -->|"loads matrices"| node_trainer
node_trainer -->|"loads batches"| node_batching
node_trainer -->|"trains and evaluates"| node_model
node_batching -->|"returns batches"| node_trainer
node_trainer -->|"requests evaluation"| node_metrics
node_metrics -->|"returns AUPRC"| node_trainer
node_trainer -->|"saves weights"| node_checkpoints
node_model -->|"computes logits"| node_prediction

click node_mimic "https://github.com/nyumedml/gnn_for_ehr/blob/master/data/mimic/preprocess_mimic.py"
click node_eicu "https://github.com/nyumedml/gnn_for_ehr/blob/master/data/eicu/preprocess_eicu.py"
click node_processed "https://github.com/nyumedml/gnn_for_ehr/tree/master/data"
click node_trainer "https://github.com/nyumedml/gnn_for_ehr/blob/master/train.py"
click node_batching "https://github.com/nyumedml/gnn_for_ehr/blob/master/utils.py"
click node_model "https://github.com/nyumedml/gnn_for_ehr/blob/master/model.py"
click node_metrics "https://github.com/nyumedml/gnn_for_ehr/blob/master/utils.py"
click node_interpretation "https://github.com/nyumedml/gnn_for_ehr/tree/master/plots"

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_mimic,node_eicu toneBlue
class node_trainer,node_batching,node_model,node_interpretation toneAmber
class node_processed,node_metrics,node_checkpoints toneMint
class node_researcher,node_rawehr,node_prediction toneIndigo
```

## Components

### EHR preparation

- [MIMIC preprocessing](https://github.com/nyumedml/gnn_for_ehr/blob/master/data/mimic/preprocess_mimic.py): `data/mimic/preprocess_mimic.py`
- [eICU preprocessing](https://github.com/nyumedml/gnn_for_ehr/blob/master/data/eicu/preprocess_eicu.py): `data/eicu/preprocess_eicu.py`

### Model learning

- [Training and evaluation loop](https://github.com/nyumedml/gnn_for_ehr/blob/master/train.py): `train.py`
- [Batch loading](https://github.com/nyumedml/gnn_for_ehr/blob/master/utils.py): `utils.py`
- [Variational graph model](https://github.com/nyumedml/gnn_for_ehr/blob/master/model.py): `model.py`
- [Singular-value analysis](https://github.com/nyumedml/gnn_for_ehr/tree/master/plots): `plots`

### Data and results

- [Processed splits](https://github.com/nyumedml/gnn_for_ehr/tree/master/data): `data`
- [Validation AUPRC](https://github.com/nyumedml/gnn_for_ehr/blob/master/utils.py): `utils.py`
- **Model checkpoints**

### Other components

- **Researcher**
- **Raw EHR records**
- **Disease-outcome predictions**

## Connections

- Researcher → MIMIC preprocessing: runs preprocessing
- Researcher → eICU preprocessing: runs preprocessing
- Raw EHR records → MIMIC preprocessing: reads records
- Raw EHR records → eICU preprocessing: reads records
- MIMIC preprocessing → Processed splits: writes splits
- eICU preprocessing → Processed splits: writes splits
- Researcher → Training and evaluation loop: starts training
- Processed splits → Training and evaluation loop: loads matrices
- Training and evaluation loop → Batch loading: loads batches
- Training and evaluation loop → Variational graph model: trains and evaluates
- Batch loading → Training and evaluation loop: returns batches
- Training and evaluation loop → Validation AUPRC: requests evaluation
- Validation AUPRC → Training and evaluation loop: returns AUPRC
- Training and evaluation loop → Model checkpoints: saves weights
- Variational graph model → Disease-outcome predictions: computes logits