pyhealth

PyHealth

Overview

PyHealth provides an end-to-end pipeline for healthcare ML on EHR data: data loading → medical code processing → patient-level dataset construction → model training → evaluation. It natively supports MIMIC-III, MIMIC-IV, eICU-CRD, and OMOP-CDM structured databases, and handles the idiosyncratic data formats of each. Medical codes (ICD-9, ICD-10, ATC, NDC, SNOMED) are organized in a hierarchical code system that supports code-level embedding and cross-ontology mapping. Pre-built tasks — mortality prediction, drug recommendation, readmission, length-of-stay, diagnosis code prediction — can be instantiated in a few lines. Custom tasks follow a standardized interface.

When to Use

• Training clinical outcome prediction models (mortality, readmission, LOS) from MIMIC-III or MIMIC-IV
• Building drug recommendation or drug interaction prediction models using ATC code hierarchy
• Processing OMOP-CDM formatted data from institutional EHR systems for ML
• Using pretrained clinical models (RETAIN, GRASP, MedBERT) as baselines on healthcare benchmarks
• Constructing patient visit sequences with temporal structure for RNN/Transformer models
• Evaluating clinical prediction models with appropriate metrics (AUROC, AUPRC, F1, Jaccard)
• Use FIDDLE for pure EHR preprocessing without ML; use clinical-longformer for clinical note NLP

Prerequisites

• Python packages: pyhealth, torch, pandas, scikit-learn
• Data requirements: MIMIC-III/IV CSV files (requires PhysioNet credentialing), eICU, or OMOP-CDM database
• MIMIC access: request at physionet.org (free; requires CITI training, ~1 week)

pip install pyhealth torch pandas scikit-learn
# Download MIMIC-III: https://physionet.org/content/mimiciii/
# Download MIMIC-IV: https://physionet.org/content/mimiciv/

Quick Start

from pyhealth.datasets import MIMIC3Dataset

# Load MIMIC-III (specify path to downloaded CSV files)
dataset = MIMIC3Dataset(
    root="path/to/mimic-iii/",
    tables=["DIAGNOSES_ICD", "PRESCRIPTIONS", "PROCEDURES_ICD"],
    code_mapping={"ICD9CM": "CCSCM"},  # map ICD-9 codes to CCS multi-level
    dev=True,  # dev=True uses 1% of data for fast testing
)

print(f"Patients: {dataset.stat()['num_patients']}")
print(f"Visits: {dataset.stat()['num_visits']}")

Core API

Module 1: Dataset Loading

Load MIMIC-III, MIMIC-IV, eICU, and OMOP-CDM datasets.

from pyhealth.datasets import MIMIC3Dataset, MIMIC4Dataset, eICUDataset

# MIMIC-III
mimic3 = MIMIC3Dataset(
    root="data/mimic-iii/",
    tables=["DIAGNOSES_ICD", "PRESCRIPTIONS", "PROCEDURES_ICD", "LABEVENTS"],
    code_mapping={"ICD9CM": "CCSCM", "NDC": "ATC3"},  # standardize codes
    dev=False,
)
stats = mimic3.stat()
print(f"MIMIC-III: {stats['num_patients']} patients, {stats['num_visits']} visits")

# MIMIC-IV
mimic4 = MIMIC4Dataset(
    root="data/mimic-iv/",
    tables=["diagnoses_icd", "prescriptions", "procedures_icd"],
    code_mapping={"ICD10CM": "CCSCM"},
    dev=True,
)

# eICU
eicu = eICUDataset(
    root="data/eicu/",
    tables=["diagnosis", "medication", "treatment"],
    dev=True,
)
print(f"eICU loaded: {eicu.stat()}")

# Explore dataset structure
patient_id = list(mimic3.patients.keys())[0]
patient = mimic3.patients[patient_id]
print(f"Patient {patient_id}: {len(patient.visits)} visits")

for visit in patient.visits[:2]:
    print(f"  Visit {visit.visit_id}:")
    print(f"    Diagnoses: {visit.get_code_list('CCSCM')[:5]}")
    print(f"    Medications: {visit.get_code_list('ATC3')[:3]}")

Module 2: Task Construction

Convert raw datasets into ML-ready task datasets.

from pyhealth.tasks import mortality_prediction_mimic3_fn
from pyhealth.datasets import SampleDataset

# Mortality prediction task
# Each sample: patient's visit history → binary mortality label
mortality_dataset = SampleDataset(
    dataset=mimic3,
    task_fn=mortality_prediction_mimic3_fn,
)

print(f"Task: mortality prediction")
print(f"Samples: {len(mortality_dataset)}")

# Inspect a sample
sample = mortality_dataset[0]
print(f"Sample keys: {list(sample.keys())}")
print(f"Conditions (ICD codes): {sample['conditions'][:3]}")
print(f"Drugs (ATC codes): {sample['drugs'][:3]}")
print(f"Label (mortality): {sample['label']}")

# Custom task: 30-day readmission prediction
def readmission_30day_fn(patient):
    """Custom task function: predict 30-day readmission after discharge."""
    samples = []
    for i, visit in enumerate(patient.visits[:-1]):
        next_visit = patient.visits[i + 1]
        # Compute days between discharge and next admission
        days_gap = (next_visit.encounter_time - visit.discharge_time).days
        label = int(days_gap <= 30)

        samples.append({
            "visit_id": visit.visit_id,
            "patient_id": patient.patient_id,
            "conditions": visit.get_code_list("CCSCM"),
            "drugs": visit.get_code_list("ATC3"),
            "procedures": visit.get_code_list("ICD9PROC"),
            "label": label,
        })
    return samples

readmission_dataset = SampleDataset(dataset=mimic3, task_fn=readmission_30day_fn)
print(f"Readmission samples: {len(readmission_dataset)}")
pos_rate = sum(s["label"] for s in readmission_dataset) / len(readmission_dataset)
print(f"Positive rate (30-day readmission): {pos_rate:.2%}")

Module 3: Medical Code Systems

Work with ICD, ATC, NDC, and other hierarchical medical code systems.

from pyhealth.medcode import InnerMap

# ICD-9-CM diagnosis codes
icd9 = InnerMap.load("ICD9CM")
code = "428.0"  # Heart failure, unspecified
print(f"Code: {code}")
print(f"Description: {icd9.lookup(code)}")
print(f"Ancestors: {icd9.get_ancestors(code)}")
print(f"Children: {icd9.get_children(code)[:5]}")

# ATC drug classification
atc = InnerMap.load("ATC")
drug_code = "A10BA02"  # Metformin
print(f"\nATC: {drug_code}")
print(f"Drug: {atc.lookup(drug_code)}")
print(f"L1 class: {atc.get_ancestors(drug_code)}")

# Cross-ontology code mapping
from pyhealth.medcode import CrossMap

# Map NDC (drug product codes) to ATC level 3
ndc_to_atc = CrossMap.load("NDC", "ATC3")
ndc_code = "0069-2587-30"  # example NDC
atc3_codes = ndc_to_atc.map(ndc_code)
print(f"NDC {ndc_code} → ATC3: {atc3_codes}")

# Map ICD-9 to ICD-10
icd9_to_icd10 = CrossMap.load("ICD9CM", "ICD10CM")
icd10 = icd9_to_icd10.map("428.0")
print(f"ICD-9 428.0 → ICD-10: {icd10}")

Module 4: Model Training

Train pre-implemented clinical ML models.

from pyhealth.models import Transformer, RETAIN
from pyhealth.datasets import split_by_patient, get_dataloader
import torch

# Train/val/test split (patient-level, no leakage)
train_ds, val_ds, test_ds = split_by_patient(mortality_dataset, [0.7, 0.1, 0.2])

train_loader = get_dataloader(train_ds, batch_size=32, shuffle=True)
val_loader   = get_dataloader(val_ds,   batch_size=64, shuffle=False)
test_loader  = get_dataloader(test_ds,  batch_size=64, shuffle=False)

print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | Test: {len(test_ds)}")

# Transformer model for EHR sequence modeling
model = Transformer(
    dataset=mortality_dataset,
    feature_keys=["conditions", "drugs", "procedures"],
    label_key="label",
    mode="binary",        # binary classification
    embedding_dim=128,
    num_heads=4,
    num_layers=2,
    dropout=0.1,
)

print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

# RETAIN: Reverse Time Attention model (interpretable clinical ML)
retain_model = RETAIN(
    dataset=mortality_dataset,
    feature_keys=["conditions", "drugs"],
    label_key="label",
    mode="binary",
    embedding_dim=64,
)

Module 5: Training and Evaluation

from pyhealth.trainer import Trainer

trainer = Trainer(
    model=model,
    metrics=["pr_auc", "roc_auc", "f1"],  # PR-AUC, ROC-AUC, F1
)

trainer.train(
    train_dataloader=train_loader,
    v