PRISM

Official implementation of the paper PRISM: Mitigating EHR Data Sparsity via Learning from Missing Feature Calibrated Prototype Patient Representations, accepted at CIKM 2024.

PRISM learns robust EHR patient representations under severe feature sparsity by combining:

• feature-isolated static and dynamic embeddings,
• missing-feature-aware attention calibration with global and local missingness,
• confidence-aware prototype patient learning with a GCN,
• adaptive fusion between individual and prototype representations.

This repository is intentionally lightweight and source-first. Code is placed directly under src/; generated datasets live under the root data/ directory.

Project Layout

.
├── configs/                 # Experiment TOML configs
├── src/
│   ├── main.py              # Thin argparse script for running stages
│   ├── config.py            # Dataclass-based TOML config loader
│   ├── datasets/            # MIMIC-IV demo conversion and PRISM tensor prep
│   ├── evaluation/          # AUROC, AUPRC, F1, calibration metrics
│   ├── models/              # PRISM model implementation
│   ├── training/            # Training loop and artifact writing
│   └── utils.py
├── tests/                   # Unit and smoke tests
└── data/                    # Ignored local raw/table/tensor data

src/datasets/ contains dataset conversion code. The root data/ folder is reserved for generated or downloaded data and is ignored by git.

Setup

PRISM is developed with uv and Python 3.12.

uv sync
source .venv/bin/activate

No package installation step is required. Use the simple argparse script in src/main.py directly after activating the environment.

Data Pipeline

Raw MIMIC-IV demo data is first converted into a reusable three-table format compatible with the OneEHR convention:

dynamic.csv: patient_id,event_time,code,value
static.csv:  patient_id,start_time,end_time,age,gender,...
label.csv:   patient_id,label_time,label_code,label_value

PRISM-specific preprocessing then reads only these three tables and constructs:

• LOCF-imputed hourly dynamic sequences,
• static age/gender tensors,
• global feature observation rates rho,
• local time-since-last-observed features tau,
• K-means prototype patient centers.

This keeps raw dataset conversion separate from PRISM-specific tensor construction.

Quickstart

Run the full MIMIC-IV demo mortality pipeline:

python src/main.py --stage all --config configs/mimic4_demo.toml

Or run each stage explicitly:

python src/main.py --stage download --config configs/mimic4_demo.toml
python src/main.py --stage convert --config configs/mimic4_demo.toml
python src/main.py --stage preprocess --config configs/mimic4_demo.toml
python src/main.py --stage train --config configs/mimic4_demo.toml --seed 0
python src/main.py --stage test --config configs/mimic4_demo.toml --seed 0

train uses the validation split for early stopping and checkpoint selection. test is a separate stage that loads the selected checkpoint and evaluates the held-out test split.

Default outputs:

data/raw/mimic-iv-demo/2.2/                 # Downloaded PhysioNet demo files
data/tables/mimic-iv-demo/mortality/        # dynamic/static/label tables
data/processed/mimic-iv-demo/mortality/     # PRISM tensors and metadata
runs/mimic4_demo_mortality/                 # Checkpoints, metrics, predictions

Implemented Components

• Feature-isolated static MLP and dynamic multi-channel GRU encoders
• Missing-aware self-attention calibration using rho and tau
• Confidence-aware patient similarity
• GCN-based prototype patient learner
• Prototype representation fusion
• Binary mortality training with validation-based early stopping
• Separate held-out test evaluation from the selected checkpoint
• AUROC, AUPRC, F1, precision, recall, Brier score
• Saved history.csv, train_val_predictions.csv, test_predictions.csv, metrics JSON files, and checkpoint artifacts

Reproduction Scope

The default config uses the public MIMIC-IV demo. It is useful for validating the full engineering pipeline, but it is too small to reproduce the paper's benchmark numbers.

To reproduce the CIKM paper tables, run the same pipeline on the full benchmark datasets used in the paper: MIMIC-III, MIMIC-IV, PhysioNet Challenge 2012, and eICU, with the reported 70/10/20 stratified split and five seeds.

Validation

uv run --extra test pytest -q
uv run --extra test ruff check .

Current smoke-test run on the MIMIC-IV demo with seed=0 completes end-to-end and writes metrics under runs/mimic4_demo_mortality/.

Citation

ACM reference:

Yinghao Zhu, Zixiang Wang, Long He, Shiyun Xie, Xiaochen Zheng, Liantao Ma, and Chengwei Pan. 2024. PRISM: Mitigating EHR Data Sparsity via Learning from Missing Feature Calibrated Prototype Patient Representations. In Proceedings of the 33rd ACM International Conference on Information and Knowledge Management (CIKM '24), 3560-3569. https://doi.org/10.1145/3627673.3679521

BibTeX:

@inproceedings{zhu2024prism,
  title = {PRISM: Mitigating EHR Data Sparsity via Learning from Missing Feature Calibrated Prototype Patient Representations},
  author = {Zhu, Yinghao and Wang, Zixiang and He, Long and Xie, Shiyun and Zheng, Xiaochen and Ma, Liantao and Pan, Chengwei},
  booktitle = {Proceedings of the 33rd ACM International Conference on Information and Knowledge Management},
  series = {CIKM '24},
  pages = {3560--3569},
  year = {2024},
  publisher = {Association for Computing Machinery},
  doi = {10.1145/3627673.3679521}
}