Claude Skill

alterlab-pyhealth

Develops, tests, and validates clinical machine learning models with the PyHealth 2.x healthcare AI toolkit. Use when working with electronic health records (EHR), clinical prediction tasks (mortality, readmission, length of stay, drug recommendation), medical coding systems (ICD

LLM Mart · 0 points · 0 0 0

Virus-scanned Reviewed automatically before listing.

Full trust report

Download alterlab-ieu-alterlab-academic-skills-skills_clinical-research_alterlab-pyhealth-e4836c0.zip · 38 KB

alterlab-ieu/alterlab-academic-skills 68 13 forks MIT Updated 1d ago

Part of alterlab-ieu/alterlab-academic-skills — 94 skills

Install

skills CLI

npx skills add https://github.com/AlterLab-IEU/AlterLab-Academic-Skills/tree/main/skills/clinical-research/alterlab-pyhealth

Claude Code

claude plugin marketplace add https://llmmart.ai/marketplace.json && claude plugin install alterlab-ieu-alterlab-academic-skills@llmmart

Git git clone https://github.com/AlterLab-IEU/AlterLab-Academic-Skills.git

The skills CLI installs just this skill, for any of its supported agents. Claude Code installs the whole alterlab-ieu/alterlab-academic-skills collection as a plugin from our marketplace. Git is the plain clone.

Skill manifest

PyHealth: Healthcare AI Toolkit

Overview

PyHealth is a Python library for healthcare AI that provides datasets, task definitions, models, trainers, and medical-code utilities for clinical machine learning. Use this skill when developing healthcare prediction models, processing clinical data, working with medical coding systems, or validating models before any clinical use.

Version gotcha (read first). This skill targets PyHealth 2.x (current release 2.0.2, Sept 2026). The 2.0 rewrite changed the API in ways most tutorials and pre-2025 snippets get wrong:

Tasks are classes you instantiate, e.g. MortalityPredictionMIMIC4(), DrugRecommendationMIMIC3() — not the old snake-case mortality_prediction_mimic4_fn functions. Pass the instance to dataset.set_task(task).

Datasets take an explicit table list. Single-source loaders use root= + tables=[...] (MIMIC3Dataset, MIMIC4EHRDataset, eICUDataset, OMOPDataset); the multimodal MIMIC4Dataset uses ehr_root= + ehr_tables=[...] (plus optional note_root/cxr_root).

Models take only the SampleDataset plus hyperparameters, e.g. Transformer(dataset=samples, embedding_dim=128). Feature keys, label key, and mode are read from the task's input_schema / output_schema; the 1.x feature_keys= / label_key= / mode= arguments raise TypeError.

Metric names have no _score suffix: pr_auc, roc_auc, f1; multilabel/drug-rec use the *_samples family (jaccard_samples, f1_samples, pr_auc_samples). Pass metrics=[...] to the Trainer constructor and monitor= one of those names.

Checkpoints use trainer.save_ckpt(path) / trainer.load_ckpt(path) (there is no trainer.save).

2.0.2 requires Python 3.12 or 3.13 and pins its own stack (numpy 2.2, pandas 2.3, torch 2.7, transformers 4.53), so install it in a dedicated environment rather than next to pandas 3 / transformers 5.

When unsure of a class or argument name, check the installed source rather than trusting older snippets.

When to Use This Skill

Invoke this skill when:

Working with healthcare datasets: MIMIC-III, MIMIC-IV, eICU, OMOP, sleep EEG data, medical images
Clinical prediction tasks: Mortality prediction, hospital readmission, length of stay, drug recommendation
Medical coding: Translating between ICD-9/10, NDC, RxNorm, ATC, CCS coding systems
Processing clinical data: Sequential events, physiological signals, clinical text, medical images
Implementing healthcare models: RETAIN, SafeDrug, GAMENet, StageNet, Transformer for EHR
Evaluating clinical models: Fairness metrics, calibration, interpretability, uncertainty quantification

Does NOT Trigger

Scenario	Use Instead
Cleaning a raw ECG/EEG/EDA trace and computing HRV or SCR features (no model training)	`alterlab-neurokit2`
Reading, anonymizing, or converting DICOM image files	`alterlab-pydicom`
Kaplan-Meier / Cox time-to-event modeling on a tabular clinical dataset	`alterlab-scikit-survival`
Biomarker-stratified cohort report with GRADE-graded treatment recommendations	`alterlab-clinical-decision`
General tabular ML on non-EHR data (scikit-learn pipelines)	`alterlab-scikit-learn`

Core Capabilities

PyHealth operates through a modular 5-stage pipeline:

Data Loading: Standardized loaders for EHR, signal, imaging, and text datasets
Task Definition: Predefined clinical prediction tasks (task classes) or custom BaseTask subclasses
Model Selection: Baselines, general deep learning, and healthcare-specific models
Training: Trainer with best-checkpoint selection, monitoring, and evaluation
Validation: Calibration, conformal prediction, fairness metrics, and interpretability methods

PyHealth 2.x uses a polars-backed data layer and caches task samples, which keeps large EHR tables memory-efficient.

Quick Start Workflow

from pyhealth.datasets import MIMIC4EHRDataset, split_by_patient, get_dataloader
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer

# 1. Load dataset (declare the tables the task needs) and set the task (a class instance)
dataset = MIMIC4EHRDataset(
    root="/path/to/mimic-iv/2.2",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())

# 2. Split data by patient (no leakage across splits)
train, val, test = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=42)

# 3. Create data loaders
train_loader = get_dataloader(train, batch_size=64, shuffle=True)
val_loader = get_dataloader(val, batch_size=64, shuffle=False)
test_loader = get_dataloader(test, batch_size=64, shuffle=False)

# 4. Initialize the model: inputs, label ("mortality"), and mode ("binary")
#    all come from the task schema, so only hyperparameters are passed
model = Transformer(dataset=sample_dataset, embedding_dim=128)

trainer = Trainer(model=model, metrics=["pr_auc", "roc_auc", "f1"])  # device auto-detected
trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    monitor="pr_auc",            # AUPRC — robust for the rare-mortality class
    monitor_criterion="max",
)

# 5. Evaluate (uses the metrics passed to the Trainer)
results = trainer.evaluate(test_loader)

Detailed Documentation

Read the reference file that matches the step you are on:

File	Read when	Key topics
`references/datasets.md`	Loading MIMIC/eICU/OMOP/signal datasets, splitting data	Patient/Event structures, loaders, `split_by_patient` / `split_by_visit` / `split_by_sample`
`references/medical_coding.md`	Translating or grouping ICD, NDC, RxNorm, ATC, CCS codes	`InnerMap` lookups and hierarchy, `CrossMap` translation
`references/tasks.md`	Choosing a predefined task or writing a custom one	2.x task classes, `input_schema` / `output_schema`, custom `BaseTask`
`references/models.md`	Selecting and configuring a model	Baselines, RNN/CNN/Transformer, RETAIN, SafeDrug, GAMENet, StageNet, GAT/GCN
`references/preprocessing.md`	Understanding how raw events become tensors	Schema-string processors (`"sequence"`, `"timeseries"`, `"binary"`, ...)
`references/training_evaluation.md`	Training, metrics, calibration, uncertainty, interpretability	`Trainer`, metric strings, conformal prediction, Chefer/IG attributions

Installation

uv venv --python 3.13 .venv-pyhealth      # PyHealth 2.0.2 supports Python 3.12-3.13
source .venv-pyhealth/bin/activate
uv pip install "pyhealth>=2.0.2"

Requirements (PyHealth 2.0.2):

Python 3.12 or 3.13 (>=3.12,<3.14) — if your default interpreter is 3.14, create the environment with --python 3.13.
PyTorch, polars, pandas, scikit-learn, and transformers are installed as pinned dependencies — keep PyHealth in its own environment so these pins don't collide with other projects.
A 2.1 alpha line exists on PyPI; stay on the 2.0.x releases unless you need an alpha-only feature.

Common Use Cases

Use Case 1: ICU Mortality Prediction

Objective: Predict patient mortality in intensive care unit

Approach:

Load MIMIC-IV dataset → Read references/datasets.md
Apply mortality prediction task → Read references/tasks.md
Select an interpretable model (RETAIN) or an attribution-friendly one (Transformer) → Read references/models.md
Train and evaluate → Read references/training_evaluation.md
Interpret predictions for clinical review → Read references/training_evaluation.md

Use Case 2: Safe Medication Recommendation

Objective: Recommend medications while avoiding drug-drug interactions

Approach:

Load EHR dataset (MIMIC-III/IV, eICU, or OMOP) → Read references/datasets.md
Apply a DrugRecommendation* task → Read references/tasks.md
Use SafeDrug or GAMENet, which build their DDI graphs from the dataset → Read references/models.md
Preprocess medication codes → Read references/medical_coding.md
Evaluate with multi-label metrics (jaccard_samples, f1_samples, pr_auc_samples) → Read references/training_evaluation.md

Use Case 3: Hospital Readmission Prediction

Objective: Identify patients at risk of readmission

Approach:

Load multi-site EHR data (eICU or OMOP) → Read references/datasets.md
Apply a ReadmissionPrediction* task → Read references/tasks.md
Handle class imbalance (report AUPRC, not only AUROC) → Read references/training_evaluation.md
Train a Transformer or RNN model → Read references/models.md
Calibrate predictions and assess fairness → Read references/training_evaluation.md

Use Case 4: Sleep Staging

Objective: Classify sleep stages from EEG signals

Approach:

Load a sleep EEG dataset (SleepEDF, SHHS, ISRUC) → Read references/datasets.md
Apply sleep staging (SleepStagingSleepEDF or the legacy sleep_staging_*_fn functions) → Read references/tasks.md
Preprocess EEG signals (filtering, segmentation) → Read references/preprocessing.md
Train a CNN, SparcNet, or ContraWR model → Read references/models.md
Evaluate per-stage performance (f1_macro, cohen_kappa) → Read references/training_evaluation.md

Use Case 5: Medical Code Translation

Objective: Standardize diagnoses across different coding systems

Approach:

Read references/medical_coding.md for comprehensive guidance
Use CrossMap to translate between ICD-9, ICD-10, and CCS
Group codes into clinically meaningful categories
Integrate with dataset processing

Use Case 6: Clinical Text to ICD Coding

Objective: Automatically assign ICD codes from clinical notes

Approach:

Load MIMIC-III with clinical notes → Read references/datasets.md
Apply the MIMIC3ICD9Coding task → Read references/tasks.md
Preprocess clinical text → Read references/preprocessing.md
Use TransformersModel(dataset=..., model_name="emilyalsentzer/Bio_ClinicalBERT") → Read references/models.md
Evaluate with multi-label metrics → Read references/training_evaluation.md

Best Practices

Data Handling

Always split by patient: Prevent data leakage by ensuring no patient appears in multiple splits

from pyhealth.datasets import split_by_patient
train, val, test = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=42)

Check dataset statistics: Understand your data before modeling

dataset.stats()  # prints patient and event counts (returns None)

Use appropriate preprocessing: Match processors to data types (see references/preprocessing.md)

Model Development

Start with baselines: Establish baseline performance with simple models
- LogisticRegression for binary/multi-class tasks
- MLP for an initial deep learning baseline
Choose task-appropriate models:
- Interpretability needed → RETAIN, AdaCare (by design); Transformer (post-hoc attributions)
- Drug recommendation → SafeDrug, GAMENet
- Long sequences → Transformer
- Graph relationships → GAT / GCN
Monitor validation metrics: Use appropriate metrics for the task and handle class imbalance. PyHealth metric strings (pass to Trainer(metrics=[...]) / monitor=):
- Binary: roc_auc, pr_auc (prefer pr_auc for rare events), f1, accuracy
- Multi-class: f1_macro, f1_weighted, accuracy, cohen_kappa
- Multi-label / drug-rec: jaccard_samples, f1_samples, pr_auc_samples, ddi (reported as ddi_score)
- Regression: mae, mse, kl_divergence

Clinical Validation

Calibrate predictions: Ensure probabilities are reliable (see references/training_evaluation.md)
Assess fairness: Evaluate across demographic groups to detect bias
Quantify uncertainty: Provide confidence estimates for predictions (conformal prediction sets)
Interpret predictions: Attention/relevance maps, SHAP, or integrated gradients for clinician review
Validate thoroughly: Use held-out test sets from different time periods or sites
Report transparently: Follow TRIPOD+AI (BMJ 2024;385:e078378) when publishing a clinical prediction model

Limitations and Considerations

Data Requirements

Large datasets: Deep learning models require sufficient data (thousands of patients)
Data quality: Missing data and coding errors impact performance
Temporal consistency: Ensure train/test split respects temporal ordering when needed
Access: MIMIC and eICU require PhysioNet credentialing and a data use agreement; never copy restricted records into prompts, notebooks, or repositories that the agreement does not cover

Clinical Validation

External validation: Test on data from different hospitals/systems
Prospective evaluation: Validate in real clinical settings before deployment
Clinical review: Have clinicians review predictions and interpretations
Decision support, not diagnosis: Present model outputs as research-grade risk estimates for qualified clinicians; deployment as a medical device falls under device regulation (e.g. FDA SaMD, EU MDR)
Ethical considerations: Address privacy (HIPAA/GDPR), fairness, and safety

Computational Resources

GPU recommended: For training deep learning models efficiently
Memory requirements: Large datasets may require 16GB+ RAM
Storage: Healthcare datasets can be 10s-100s of GB, plus the task-sample cache

Troubleshooting

Common Issues

TypeError: ... unexpected keyword argument 'feature_keys' (or 'root'):

You are using a 1.x-style call. Pass only dataset= and hyperparameters to models; use MIMIC4EHRDataset(root=..., tables=...) or MIMIC4Dataset(ehr_root=..., ehr_tables=...)

ImportError or missing tables:

Ensure dataset files are downloaded and the root path points at the versioned folder
Confirm the table names exist in the dataset's YAML config

Out of memory:

Reduce batch size
Reduce sequence length (max_seq_len on Transformer)
Pass dev=True to the dataset loader (e.g. MIMIC4EHRDataset(..., dev=True)) to prototype on the first 1,000 patients
Process data in chunks

Poor performance:

Check class imbalance and use appropriate metrics (pr_auc vs roc_auc)
Verify preprocessing (normalization, missing data handling)
Increase model capacity or training epochs
Check for data leakage in the train/test split

Slow training:

Use a GPU (Trainer(..., device="cuda"))
Increase batch size (if memory allows)
Reduce sequence length
Use a lighter model (CNN or RNN instead of Transformer)

Getting Help

Documentation: https://pyhealth.readthedocs.io/
GitHub Issues: https://github.com/sunlabuiuc/PyHealth/issues
Examples/notebooks: https://github.com/sunlabuiuc/PyHealth/tree/master/examples

Example: Complete Workflow

# Complete mortality prediction pipeline (PyHealth 2.0.x)
import torch
from pyhealth.datasets import MIMIC4EHRDataset, split_by_patient, get_dataloader
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer
from pyhealth.interpret.methods import CheferRelevance

# 1. Load dataset (declare the tables the task needs)
dataset = MIMIC4EHRDataset(
    root="/data/mimic-iv/2.2",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
dataset.stats()

# 2. Define task (instantiate the task class)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())
print(f"Generated {len(sample_dataset)} samples")

# 3. Split data (by patient to prevent leakage)
train_ds, val_ds, test_ds = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=42)

# 4. Create data loaders
train_loader = get_dataloader(train_ds, batch_size=64, shuffle=True)
val_loader = get_dataloader(val_ds, batch_size=64)
test_loader = get_dataloader(test_ds, batch_size=64)

# 5. Initialize the model (schema-driven; swap in RETAIN(dataset=sample_dataset,
#    embedding_dim=128) for a model that is interpretable by design)
model = Transformer(dataset=sample_dataset, embedding_dim=128, heads=2, num_layers=2)

# 6. Train, keeping the best checkpoint by validation AUPRC
trainer = Trainer(model=model, metrics=["accuracy", "pr_auc", "roc_auc", "f1"])
trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    optimizer_class=torch.optim.Adam,
    optimizer_params={"lr": 1e-3},
    weight_decay=1e-5,
    monitor="pr_auc",          # AUPRC for the imbalanced (rare-mortality) outcome
    monitor_criterion="max",
    patience=5,                # early stopping
)

# 7. Evaluate on the test set (uses the metrics passed to the Trainer)
for metric, value in trainer.evaluate(test_loader).items():
    print(f"  {metric}: {value:.4f}")

# 8. Predictions with patient IDs: inference() returns (y_true, y_prob, loss),
#    extended with patient_ids when return_patient_ids=True
y_true, y_prob, loss, patient_ids = trainer.inference(test_loader, return_patient_ids=True)
positive_prob = y_prob if y_prob.ndim == 1 else y_prob[..., -1]
high_risk_idx = int(positive_prob.argmax())
print(f"Highest-risk patient: {patient_ids[high_risk_idx]} ({float(positive_prob[high_risk_idx]):.3f})")

# 9. Token-level relevance (Chefer; supported by Transformer and StageAttentionNet)
relevance = CheferRelevance(model)
batch = next(iter(get_dataloader(test_ds, batch_size=1, shuffle=False)))
for feature_key, rel in relevance.attribute(**batch).items():
    print(f"{feature_key}: top tokens -> {rel[0].topk(min(5, rel.shape[-1])).indices.tolist()}")

# 10. Save the trained weights
trainer.save_ckpt("./models/mortality_transformer.pt")

Resources

For detailed information on each component, see the reference files in references/: datasets.md, medical_coding.md, tasks.md, models.md, preprocessing.md, and training_evaluation.md (see the table under Detailed Documentation for when to read each).

Part of the AlterLab Academic Skills suite.

Files (alterlab-academic-skills)

evals

evals.json 6.2 KB

{
  "skill": "alterlab-pyhealth",
  "evals": [
    {
      "id": "mortality-prediction-mimic4",
      "prompt": "I want to predict ICU mortality on MIMIC-IV. Load the dataset, set up the mortality task, split by patient, and train an interpretable RETAIN model monitoring AUPRC since the outcome is rare.",
      "expected_output": "Invokes alterlab-pyhealth. Loads MIMIC-IV with MIMIC4EHRDataset(root=..., tables=[...]) (or MIMIC4Dataset(ehr_root=..., ehr_tables=[...])), sets the MortalityPredictionMIMIC4() task class, uses split_by_patient to prevent leakage, builds get_dataloader loaders, initializes RETAIN(dataset=sample_dataset, ...) whose inputs, 'mortality' label, and binary mode come from the task schema (no feature_keys/label_key/mode arguments in PyHealth 2.x), trains with Trainer(metrics=[...]) monitoring pr_auc (AUPRC for imbalance), and evaluates. Recommends RETAIN for interpretability.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "split_by_patient" },
        { "type": "behavior", "value": "Uses a MIMIC-IV dataset loader, a mortality task class instance, split_by_patient, RETAIN built from the sample dataset (no 1.x feature_keys/label_key arguments), and the Trainer monitoring AUPRC (pr_auc)." }
      ]
    },
    {
      "id": "safedrug-recommendation",
      "prompt": "I'm building a medication recommendation model that avoids drug-drug interactions on MIMIC-IV. Which model should I use and how do I handle the multi-label medication codes?",
      "expected_output": "Invokes alterlab-pyhealth. Recommends the SafeDrug model (or GAMENet) with DDI constraints for the drug recommendation task, frames it as multi-label, processes medication codes (NDC/RxNorm/ATC) potentially via InnerMap/CrossMap, and evaluates with multi-label metrics like Jaccard and example-F1.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "SafeDrug" },
        { "type": "behavior", "value": "Recommends SafeDrug (or GAMENet) for DDI-aware drug recommendation and uses multi-label metrics (Jaccard / example-F1)." }
      ]
    },
    {
      "id": "medical-code-crossmap",
      "prompt": "I have diagnoses coded in ICD-9 in one dataset and ICD-10 in another. I want to translate them to a common system and group them into clinically meaningful categories like CCS.",
      "expected_output": "Invokes alterlab-pyhealth medical-coding capability. Uses CrossMap to translate between ICD-9, ICD-10, and CCS, and InnerMap for within-system hierarchy lookups, grouping codes into clinical categories. References references/medical_coding.md.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "CrossMap" },
        { "type": "behavior", "value": "Uses CrossMap (cross-system) and/or InnerMap to translate ICD-9/ICD-10/CCS codes." }
      ]
    },
    {
      "id": "fairness-calibration-deployment",
      "prompt": "My readmission model on eICU performs well overall but I need to check it for bias across demographic groups, calibrate the probabilities, and quantify prediction uncertainty before any clinical deployment.",
      "expected_output": "Invokes alterlab-pyhealth training/evaluation capability. Computes fairness metrics across demographics, applies calibration (Platt or temperature scaling), and quantifies uncertainty (conformal prediction or MC dropout), framing all of this as pre-deployment clinical validation. References references/training_evaluation.md.",
      "assertions": [
        { "type": "should_trigger", "value": true },
        { "type": "output_contains", "value": "calibrat" },
        { "type": "behavior", "value": "Addresses fairness metrics, calibration (Platt/temperature scaling), and uncertainty quantification for clinical deployment." }
      ]
    },
    {
      "id": "near-miss-neurokit2",
      "prompt": "I have a single raw ECG trace at 1000 Hz and I just want to clean it, detect the R-peaks, and compute SDNN and RMSSD heart rate variability metrics from it.",
      "expected_output": "Does NOT invoke this skill; defers to alterlab-neurokit2. The task is raw physiological signal processing and HRV feature extraction from one ECG trace (nk.ecg_process, nk.hrv), not training a clinical prediction model on a healthcare dataset. PyHealth consumes datasets and trains models; it does not do low-level biosignal cleaning and HRV computation.",
      "assertions": [
        { "type": "should_not_trigger", "value": true },
        { "type": "output_contains", "value": "alterlab-neurokit2" }
      ]
    },
    {
      "id": "near-miss-clinical-decision",
      "prompt": "Analyze a cohort of 60 HER2-positive metastatic breast cancer patients stratified by hormone receptor status, with hazard ratios, Kaplan-Meier survival curves, and GRADE-graded treatment recommendations as a publication-ready PDF.",
      "expected_output": "Does NOT invoke this skill; defers to alterlab-clinical-decision. The task is a biomarker-stratified clinical decision support document (cohort survival analysis, hazard ratios, GRADE grading, LaTeX/PDF), not a machine-learning prediction pipeline on an EHR dataset. PyHealth builds and trains predictive models; it does not generate GRADE-graded cohort analysis documents.",
      "assertions": [
        { "type": "should_not_trigger", "value": true },
        { "type": "output_contains", "value": "alterlab-clinical-decision" }
      ]
    },
    {
      "id": "near-miss-scikit-survival",
      "prompt": "I have a CSV of 800 colorectal cancer patients with time-to-recurrence, a censoring flag, age, stage, and treatment arm. Fit a Cox proportional hazards model and a random survival forest and compare them by concordance index.",
      "expected_output": "Does NOT invoke this skill; defers to alterlab-scikit-survival. The task is classical time-to-event modeling on a flat tabular dataset (CoxPHSurvivalAnalysis, RandomSurvivalForest, concordance index), not a PyHealth EHR pipeline with dataset loaders, task classes, and deep clinical models.",
      "assertions": [
        { "type": "should_not_trigger", "value": true },
        { "type": "output_contains", "value": "alterlab-scikit-survival" }
      ]
    }
  ]
}

references

datasets.md 6.3 KB

# PyHealth Datasets and Data Structures

## Core Data Structures

PyHealth 2.x stores every source table as rows of one long, polars-backed event table; the 1.x `Visit` object is deprecated.

### Event
One row from a source table:
- **event_type**: The table it came from (e.g. `"admissions"`, `"diagnoses_icd"`, `"prescriptions"`)
- **timestamp**: `datetime` of the event
- **attributes**: The table's configured columns, readable as attributes (e.g. `event.icd_code`, `event.icd_version`, `event.hadm_id`, `event.dischtime` for MIMIC-IV); column names come from the dataset YAML config

### Patient
All events for one patient, sorted by time:
- **patient_id**: Unique identifier
- **get_events(event_type=..., start=..., end=...)**: Events of one type, optionally within a time window — tasks use this to collect, e.g., the diagnoses recorded during an admission
- Visits/admissions are just events of type `"admissions"` (or the dataset's equivalent)

## BaseDataset Class

In PyHealth 2.x, dataset constructors take an explicit table list naming the source tables to load (table names come from the dataset's YAML config, e.g. `mimic4_ehr.yaml`). Single-source loaders (`MIMIC3Dataset`, `MIMIC4EHRDataset`, `eICUDataset`, `OMOPDataset`) take `root=` + `tables=[...]`; the multimodal `MIMIC4Dataset` takes `ehr_root=` / `note_root=` / `cxr_root=` with `ehr_tables=` / `note_tables=` / `cxr_tables=`. The EHR loaders (MIMIC-III/IV, eICU, OMOP) accept `dev=True` to work on the first 1,000 patients.

**Key Methods:**
- `iter_patients()`: Iterate through all patients
- `stats()`: Print dataset statistics (patient and event counts; returns `None`)
- `set_task(task)`: Apply a prediction task. Pass an **instance** of a task class (e.g. `MortalityPredictionMIMIC4()`), not a bare function.

## Available Datasets

### Electronic Health Record (EHR) Datasets

**MIMIC-III Dataset** (`MIMIC3Dataset`)
- Intensive care unit data from Beth Israel Deaconess Medical Center
- 40,000+ critical care patients
- Diagnoses, procedures, medications, lab results
- Usage: `from pyhealth.datasets import MIMIC3Dataset`

**MIMIC-IV Datasets** (`MIMIC4EHRDataset`, `MIMIC4NoteDataset`, `MIMIC4CXRDataset`, `MIMIC4Dataset`)
- Hospital + ICU EHR from Beth Israel Deaconess (MIMIC-IV), with optional notes and chest X-rays
- `MIMIC4EHRDataset(root=..., tables=[...])` for structured EHR only
- `MIMIC4Dataset(ehr_root=..., ehr_tables=[...], note_root=..., cxr_root=...)` to combine modalities
- Usage: `from pyhealth.datasets import MIMIC4EHRDataset`

**eICU Dataset** (`eICUDataset`)
- Multi-center critical care database
- 200,000+ admissions from 200+ hospitals
- Standardized ICU data across facilities
- Usage: `from pyhealth.datasets import eICUDataset`

**OMOP Dataset** (`OMOPDataset`)
- Observational Medical Outcomes Partnership format
- Standardized common data model
- Interoperability across healthcare systems
- Usage: `from pyhealth.datasets import OMOPDataset`

**EHRShot Dataset** (`EHRShotDataset`)
- Benchmark dataset for few-shot learning
- Specialized for testing model generalization
- Usage: `from pyhealth.datasets import EHRShotDataset`

### Physiological Signal Datasets

**Sleep EEG Datasets:**
- `SleepEDFDataset`: Sleep-EDF database for sleep staging
- `SHHSDataset`: Sleep Heart Health Study data
- `ISRUCDataset`: ISRUC-Sleep database

**Temple University EEG Datasets:**
- `TUEVDataset`: Abnormal EEG events detection
- `TUABDataset`: Abnormal/normal EEG classification
- `TUSZDataset`: Seizure detection

**All signal datasets support:**
- Multi-channel EEG signals
- Standardized sampling rates
- Expert annotations
- Sleep stage or abnormality labels

### Medical Imaging Datasets

**COVID-19 CXR Dataset** (`COVID19CXRDataset`)
- Chest X-ray images for COVID-19 classification
- Multi-class labels (COVID-19, pneumonia, normal)
- Usage: `from pyhealth.datasets import COVID19CXRDataset`

### Text-Based Datasets

**Medical Transcriptions Dataset** (`MedicalTranscriptionsDataset`)
- Clinical notes and transcriptions
- Medical specialty classification
- Text-based prediction tasks
- Usage: `from pyhealth.datasets import MedicalTranscriptionsDataset`

**Cardiology Dataset** (`CardiologyDataset`)
- Cardiac patient records
- Cardiovascular disease prediction
- Usage: `from pyhealth.datasets import CardiologyDataset`

### Preprocessed Datasets

**MIMIC Extract Dataset** (`MIMICExtractDataset`)
- Pre-extracted MIMIC features
- Ready-to-use benchmarking data
- Reduced preprocessing requirements
- Usage: `from pyhealth.datasets import MIMICExtractDataset`

## SampleDataset Class

Converts raw datasets into task-specific formatted samples.

**Purpose:** Transform patient-level data into model-ready input/output pairs

**Key Attributes:**
- `input_schema`: Defines input data structure
- `output_schema`: Defines target labels/predictions
- `samples`: List of processed samples

**Usage Pattern:**
```python
# After setting a task (a task-class instance) on a BaseDataset
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())
```

## Data Splitting Functions

**Patient-Level Split** (`split_by_patient`)
- Ensures no patient appears in multiple splits
- Prevents data leakage
- Recommended for clinical prediction tasks

**Visit-Level Split** (`split_by_visit`)
- Splits by individual visits
- Allows same patient across splits (use cautiously)

**Sample-Level Split** (`split_by_sample`)
- Random sample splitting
- Most flexible but may cause leakage

**Parameters:**
- `dataset`: SampleDataset to split
- `ratios`: Tuple of split ratios (e.g., [0.7, 0.1, 0.2])
- `seed`: Random seed for reproducibility

## Common Workflow

```python
from pyhealth.datasets import MIMIC4EHRDataset, split_by_patient
from pyhealth.tasks import MortalityPredictionMIMIC4

# 1. Load dataset (declare the tables you need)
dataset = MIMIC4EHRDataset(
    root="/path/to/data",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)

# 2. Set prediction task (instantiate the task class)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())

# 3. Split data
train, val, test = split_by_patient(sample_dataset, [0.7, 0.1, 0.2])

# 4. Get statistics (prints; returns None)
dataset.stats()
```

## Performance Notes

- PyHealth 2.x uses a **polars-backed** data layer for fast, memory-efficient processing
- Optimized for large-scale EHR datasets
- Memory-efficient patient iteration
- Vectorized operations for feature extraction

medical_coding.md 9 KB

# PyHealth Medical Code Translation

## Overview

Healthcare data uses multiple coding systems and standards. PyHealth's MedCode module enables translation and mapping between medical coding systems through ontology lookups and cross-system mappings.

## Core Classes

### InnerMap
Handles within-system ontology lookups and hierarchical navigation.

**Key Capabilities:**
- Code lookup with attributes (names, descriptions)
- Ancestor/descendant hierarchy traversal
- Code standardization and conversion
- Parent-child relationship navigation

### CrossMap
Manages cross-system mappings between different coding standards.

**Key Capabilities:**
- Translation between coding systems
- Many-to-many relationship handling
- Hierarchical level specification (for medications)
- Bidirectional mapping support

## Supported Coding Systems

### Diagnosis Codes

**ICD-9-CM (International Classification of Diseases, 9th Revision, Clinical Modification)**
- Legacy diagnosis coding system
- Hierarchical structure with 3-5 digit codes
- Used in US healthcare pre-2015
- Usage: `from pyhealth.medcode import InnerMap`
  - `icd9_map = InnerMap.load("ICD9CM")`

**ICD-10-CM (International Classification of Diseases, 10th Revision, Clinical Modification)**
- Current diagnosis coding standard
- Alphanumeric codes (3-7 characters)
- More granular than ICD-9
- Usage: `from pyhealth.medcode import InnerMap`
  - `icd10_map = InnerMap.load("ICD10CM")`

**CCSCM (Clinical Classifications Software for ICD-CM)**
- Groups ICD codes into clinically meaningful categories
- Reduces dimensionality for analysis
- Single-level and multi-level hierarchies
- Usage: `from pyhealth.medcode import CrossMap`
  - `icd_to_ccs = CrossMap.load("ICD9CM", "CCSCM")`

### Procedure Codes

**ICD-9-PROC (ICD-9 Procedure Codes)**
- Inpatient procedure classification
- 3-4 digit numeric codes
- Legacy system (pre-2015)
- Usage: `from pyhealth.medcode import InnerMap`
  - `icd9proc_map = InnerMap.load("ICD9PROC")`

**ICD-10-PROC (ICD-10 Procedure Coding System)**
- Current procedural coding standard
- 7-character alphanumeric codes
- More detailed than ICD-9-PROC
- Usage: `from pyhealth.medcode import InnerMap`
  - `icd10proc_map = InnerMap.load("ICD10PROC")`

**CCSPROC (Clinical Classifications Software for Procedures)**
- Groups procedure codes into categories
- Simplifies procedure analysis
- Usage: `from pyhealth.medcode import CrossMap`
  - `proc_to_ccs = CrossMap.load("ICD9PROC", "CCSPROC")`

### Medication Codes

**NDC (National Drug Code)**
- US FDA drug identification system
- 10 or 11-digit codes
- Product-level specificity (manufacturer, strength, package)
- Usage: `from pyhealth.medcode import InnerMap`
  - `ndc_map = InnerMap.load("NDC")`

**RxNorm**
- Standardized drug terminology
- Normalized drug names and relationships
- Links multiple drug vocabularies
- Usage: `from pyhealth.medcode import CrossMap`
  - `ndc_to_rxnorm = CrossMap.load("NDC", "RXNORM")`

**ATC (Anatomical Therapeutic Chemical Classification)**
- WHO drug classification system
- 5-level hierarchy:
  - **Level 1**: Anatomical main group (1 letter)
  - **Level 2**: Therapeutic subgroup (2 digits)
  - **Level 3**: Pharmacological subgroup (1 letter)
  - **Level 4**: Chemical subgroup (1 letter)
  - **Level 5**: Chemical substance (2 digits)
- Example: "C03CA01" = Furosemide
  - C = Cardiovascular system
  - C03 = Diuretics
  - C03C = High-ceiling diuretics
  - C03CA = Sulfonamides
  - C03CA01 = Furosemide

**Usage:**
```python
from pyhealth.medcode import CrossMap
ndc_to_atc = CrossMap.load("NDC", "ATC")
# Pass target-system options via target_kwargs (NOT a bare level= arg)
atc_codes = ndc_to_atc.map("50580049698", target_kwargs={"level": 3})  # ATC level 3
```

## Common Operations

### InnerMap Operations

**1. Code Lookup**
```python
from pyhealth.medcode import InnerMap

icd9_map = InnerMap.load("ICD9CM")
info = icd9_map.lookup("428.0")  # Heart failure
# Returns: name, description, additional attributes
```

**2. Ancestor Traversal**
```python
# Get all parent codes in hierarchy
ancestors = icd9_map.get_ancestors("428.0")
# Returns: ["428", "420-429", "390-459"]
```

**3. Descendant Traversal**
```python
# Get all child codes
descendants = icd9_map.get_descendants("428")
# Returns: ["428.0", "428.1", "428.2", ...]
```

**4. Code Standardization**
```python
# Normalize code format
standard_code = icd9_map.standardize("4280")  # Returns "428.0"
```

### CrossMap Operations

**1. Direct Translation**
```python
from pyhealth.medcode import CrossMap

# ICD-9-CM to CCS
icd_to_ccs = CrossMap.load("ICD9CM", "CCSCM")
ccs_codes = icd_to_ccs.map("82101")  # Coronary atherosclerosis
# Returns: ["101"]  # CCS category for coronary atherosclerosis
```

**2. Hierarchical Drug Mapping**
```python
# NDC to ATC at different levels (level goes in target_kwargs)
ndc_to_atc = CrossMap.load("NDC", "ATC")

atc_level_1 = ndc_to_atc.map("50580049698", target_kwargs={"level": 1})  # Anatomical group
atc_level_3 = ndc_to_atc.map("50580049698", target_kwargs={"level": 3})  # Pharmacological
atc_level_5 = ndc_to_atc.map("50580049698", target_kwargs={"level": 5})  # Chemical substance
```

**3. Bidirectional Mapping**
```python
# Map in either direction
rxnorm_to_ndc = CrossMap.load("RXNORM", "NDC")
ndc_codes = rxnorm_to_ndc.map("197381")  # Get all NDC codes for RxNorm
```

## Workflow Examples

### Example 1: Standardize and Group Diagnoses
```python
from pyhealth.medcode import InnerMap, CrossMap

# Load maps
icd9_map = InnerMap.load("ICD9CM")
icd_to_ccs = CrossMap.load("ICD9CM", "CCSCM")

# Process diagnosis codes
raw_codes = ["4280", "428.0", "42800"]

standardized = [icd9_map.standardize(code) for code in raw_codes]
# All become "428.0"

ccs_categories = [icd_to_ccs.map(code)[0] for code in standardized]
# All map to CCS category "108" (Heart failure)
```

### Example 2: Drug Classification Analysis
```python
from pyhealth.medcode import CrossMap

# Map NDC to ATC for drug class analysis
ndc_to_atc = CrossMap.load("NDC", "ATC")

patient_drugs = ["00074-3799-13", "00074-7286-01", "00456-0765-01"]

# Get therapeutic subgroups (ATC level 2)
drug_classes = []
for ndc in patient_drugs:
    atc_codes = ndc_to_atc.map(ndc, target_kwargs={"level": 2})
    if atc_codes:
        drug_classes.append(atc_codes[0])

# Analyze drug class distribution
```

### Example 3: ICD-9 to ICD-10 Migration
```python
from pyhealth.medcode import CrossMap

# Load ICD-9 to ICD-10 mapping
icd9_to_icd10 = CrossMap.load("ICD9CM", "ICD10CM")

# Convert historical ICD-9 codes
icd9_code = "428.0"
icd10_codes = icd9_to_icd10.map(icd9_code)
# Returns: ["I50.9", "I50.1", ...]  # Multiple possible ICD-10 codes

# Handle one-to-many mappings
for icd10_code in icd10_codes:
    print(f"ICD-9 {icd9_code} -> ICD-10 {icd10_code}")
```

## Integration with Datasets

Medical code translation integrates seamlessly with PyHealth datasets:

The `medcode` maps are independent of the dataset classes — load them directly and apply them to codes you pull from a dataset (event-traversal attributes vary by dataset/version; confirm against your loaded tables).

```python
from pyhealth.datasets import MIMIC4EHRDataset
from pyhealth.medcode import CrossMap

# Load dataset (declare the tables you need)
dataset = MIMIC4EHRDataset(
    root="/path/to/data",
    tables=["diagnoses_icd"],
)

# Load code mapping
icd_to_ccs = CrossMap.load("ICD10CM", "CCSCM")

# Process patient diagnoses (PyHealth 2.x: events are table rows; MIMIC-IV
# diagnoses_icd rows carry icd_code and icd_version columns)
for patient in dataset.iter_patients():
    for event in patient.get_events(event_type="diagnoses_icd"):
        if str(event.icd_version) != "10":
            continue
        try:
            ccs_codes = icd_to_ccs.map(event.icd_code)  # map() standardizes the code
        except KeyError:
            ccs_codes = []                               # no CCS mapping for this code
        print(f"Diagnosis {event.icd_code} -> CCS {ccs_codes}")
```

## Use Cases

### Clinical Research
- Standardize diagnoses across different coding systems
- Group related conditions for cohort identification
- Harmonize multi-site studies with different standards

### Drug Safety Analysis
- Classify medications by therapeutic class
- Identify drug-drug interactions at class level
- Analyze polypharmacy patterns

### Healthcare Analytics
- Reduce diagnosis/procedure dimensionality
- Create meaningful clinical categories
- Enable longitudinal analysis across coding system changes

### Machine Learning
- Create consistent feature representations
- Handle vocabulary mismatch in training/test data
- Generate hierarchical embeddings

## Best Practices

1. **Always standardize codes** before mapping to ensure consistent format
2. **Handle one-to-many mappings** appropriately (some codes map to multiple targets)
3. **Specify ATC level** explicitly when mapping drugs to avoid ambiguity
4. **Use CCS categories** to reduce diagnosis/procedure dimensionality
5. **Validate mappings** as some codes may not have direct translations
6. **Document code versions** (ICD-9 vs ICD-10) to maintain data provenance

models.md 15.8 KB

# PyHealth Models

## Overview

PyHealth provides 33+ models for healthcare prediction tasks, ranging from simple baselines to state-of-the-art deep learning architectures. Models are organized into general-purpose architectures and healthcare-specific models.

## Model Base Class

All models inherit from `BaseModel` (a `torch.nn.Module`).

> **Init args (PyHealth 2.x, verified on 2.0.2).** Models take the `SampleDataset` returned by `set_task()` plus hyperparameters — e.g. `Transformer(dataset=sample_dataset, embedding_dim=128)`. `BaseModel.__init__(dataset)` reads the feature keys from the task's `input_schema`, the label key from its `output_schema`, and the mode (`"binary"`, `"multiclass"`, `"multilabel"`, `"regression"`) from the output processor. The 1.x arguments `feature_keys=`, `label_key=`, and `mode=` are **not** accepted by the core EHR models (they raise `TypeError`); only a few legacy generative/vision classes (`VAE`, `Graph_TorchvisionModel`) still take them. To change which inputs a model sees, change the task schema.

**Key Attributes (set from the dataset):**
- `dataset`: Associated SampleDataset
- `feature_keys`: List of input keys (from the task's `input_schema`)
- `label_keys`: List of label keys (from the task's `output_schema`; most models require exactly one)
- `mode`: Task type resolved from the output schema
- `embedding_dim`: Feature embedding dimension (constructor argument)

**Key Methods:**
- `forward(**batch)`: Returns a dict with `loss`, `y_prob`, `y_true`, `logit` (plus `embed` when called with `embed=True`)
- `get_output_size()`, `get_loss_function()`, `prepare_y_prob()`: helpers used by subclasses
- Checkpointing lives on the `Trainer` (`save_ckpt()` / `load_ckpt()`), not on the model

## General-Purpose Models

### Baseline Models

**Logistic Regression** (`LogisticRegression`)
- Linear classifier with mean pooling
- Simple baseline for comparison
- Fast training and inference
- Good for interpretability

**Usage:**
```python
from pyhealth.models import LogisticRegression

model = LogisticRegression(
    dataset=sample_dataset,
)
```

**Multi-Layer Perceptron** (`MLP`)
- Feedforward neural network
- Configurable hidden layers
- Mean/sum pooling of nested sequence inputs
- Good baseline for structured data

**Parameters (2.0.2):**
- `embedding_dim`: Embedding size (default 128)
- `hidden_dim`: Hidden layer size (default 128)
- `n_layers`: Number of MLP layers (default 2)
- `activation`: Activation name (default `"relu"`)
- Nested sequence inputs are mean/sum-pooled before the MLP

**Usage:**
```python
from pyhealth.models import MLP

model = MLP(
    dataset=sample_dataset,
    hidden_dim=128,
)
```

### Convolutional Neural Networks

**CNN** (`CNN`)
- Convolutional layers for pattern detection
- Effective for sequential and spatial data
- Captures local temporal patterns
- Parameter efficient

**Architecture:**
- Multiple 1D convolutional layers
- Max pooling for dimension reduction
- Fully connected output layers

**Parameters (2.0.2):**
- `embedding_dim`: Embedding size (default 128)
- `hidden_dim`: Convolution channels (default 128)
- `num_layers`: Number of conv layers (default 1)

**Usage:**
```python
from pyhealth.models import CNN

model = CNN(
    dataset=sample_dataset,
)
```

**Temporal Convolutional Networks** (`TCN`)
- Dilated convolutions for long-range dependencies
- Causal convolutions (no future information leakage)
- Efficient for long sequences
- Good for time-series prediction

**Advantages:**
- Captures long-term dependencies
- Parallelizable (faster than RNNs)
- Stable gradients

**Usage:** `TCN(dataset=sample_dataset, embedding_dim=128, num_channels=128)` (`num_channels` may be a list, one entry per level)

### Recurrent Neural Networks

**RNN** (`RNN`)
- Basic recurrent architecture
- Supports LSTM, GRU, RNN variants
- Sequential processing
- Captures temporal dependencies

**Parameters (2.0.2):**
- `embedding_dim`, `hidden_dim`: Embedding and hidden-state sizes (default 128)
- Passed through to `RNNLayer`: `rnn_type` (`"GRU"` default, `"LSTM"`, `"RNN"`), `num_layers` (1), `dropout` (0.5), `bidirectional` (False)

**Usage:**
```python
from pyhealth.models import RNN

model = RNN(
    dataset=sample_dataset,
    rnn_type="LSTM",
    hidden_dim=128,
)
```

**Best for:**
- Sequential clinical events
- Temporal pattern learning
- Variable-length sequences

### Transformer Models

**Transformer** (`Transformer`)
- Self-attention mechanism
- Parallel processing of sequences
- State-of-the-art performance
- Effective for long-range dependencies

**Architecture:**
- Multi-head self-attention
- Position embeddings
- Feed-forward networks
- Layer normalization

**Parameters (2.0.2):**
- `embedding_dim`: Model width (default 128)
- `heads`: Attention heads per block (default 1)
- `num_layers`: Transformer blocks per feature stream (default 1)
- `dropout`: Dropout rate (default 0.5)
- `max_seq_len`: Maximum sequence length (default 1024)

**Usage:**
```python
from pyhealth.models import Transformer

model = Transformer(
    dataset=sample_dataset,
    embedding_dim=128,
    heads=2,
    num_layers=2,
    dropout=0.1,
)
```

`Transformer` implements the Chefer-relevance interface and `forward_from_embedding`, so it works with `CheferRelevance`, `AttentionRollout`, and the embedding-gradient methods (`IntegratedGradients`, `DeepLift`) in `pyhealth.interpret.methods` (see `references/training_evaluation.md`).

**TransformersModel** (`TransformersModel`)
- Integration with HuggingFace transformers
- Pre-trained language models for clinical text
- Fine-tuning for healthcare tasks
- Examples: BERT, RoBERTa, BioClinicalBERT

**Usage:**
```python
from pyhealth.models import TransformersModel

# Signature: TransformersModel(dataset, model_name, dropout=0.1)
model = TransformersModel(
    dataset=sample_dataset,          # task with a single text input and one label
    model_name="emilyalsentzer/Bio_ClinicalBERT",
)
```

### Graph Neural Networks

**GAT / GCN** (`GAT`, `GCN`) — there is no single `GNN` class in 2.x
- Graph-based learning over the task's inputs
- Models relationships between entities
- `GAT` (graph attention) and `GCN` (graph convolution) are separate classes

**Use Cases:**
- Drug-drug interactions
- Patient similarity networks
- Knowledge graph integration
- Comorbidity relationships

**Parameters (2.0.2):**
- `embedding_dim`: Embedding size (default 128)
- `nhid`: Hidden units per graph layer (default 64)
- `num_layers`: Number of graph layers (default 2)
- `dropout`: Dropout rate (default 0.5)
- `nheads`: Attention heads (`GAT` only, default 1)

**Usage:**
```python
from pyhealth.models import GAT, GCN

model = GAT(dataset=sample_dataset, embedding_dim=128, nhid=64, nheads=2)
# or: GCN(dataset=sample_dataset, embedding_dim=128, nhid=64)
```

## Healthcare-Specific Models

### Interpretable Clinical Models

**RETAIN** (`RETAIN`)
- Reverse time attention mechanism
- Highly interpretable predictions
- Visit-level and event-level attention
- Identifies influential clinical events

**Key Features:**
- Two-level attention (visits and features)
- Temporal decay modeling
- Clinically meaningful explanations
- Published in NeurIPS 2016

**Usage:**
```python
from pyhealth.models import RETAIN

model = RETAIN(
    dataset=sample_dataset,
)

# A forward pass returns a dict: loss, y_prob, y_true, logit.
# RETAIN is interpretable by design, but PyHealth does not return its
# alpha/beta attention weights and it does not implement the Chefer or
# forward_from_embedding interfaces. For post-hoc token attributions,
# train a Transformer (see references/training_evaluation.md).
out = model(**batch)
loss, y_prob = out["loss"], out["y_prob"]
```

**Best for:**
- Mortality prediction
- Readmission prediction
- Clinical risk scoring
- Interpretable predictions

**AdaCare** (`AdaCare`)
- Adaptive care model with feature calibration
- Disease-specific attention
- Handles irregular time intervals
- Interpretable feature importance

**ConCare** (`ConCare`)
- Cross-visit convolutional attention
- Temporal convolutional feature extraction
- Multi-level attention mechanism
- Good for longitudinal EHR modeling

### Medication Recommendation Models

**GAMENet** (`GAMENet`)
- Graph-based medication recommendation
- Drug-drug interaction modeling
- Memory network for patient history
- Multi-hop reasoning

**Architecture:**
- Drug knowledge graph
- Memory-augmented neural network
- DDI-aware prediction

**Usage:**
```python
from pyhealth.models import GAMENet

# GAMENet builds its EHR co-occurrence and DDI adjacency matrices from the
# dataset itself; passing ehr_adj/ddi_adj raises ValueError.
model = GAMENet(
    dataset=sample_dataset,   # from a DrugRecommendation* task
    embedding_dim=64,
    hidden_dim=64,
)
```

Reference: Shang et al., GAMENet: Graph Augmented MEmory Networks for Recommending Medication Combination, AAAI 2019.

**MICRON** (`MICRON`)
- Medication recommendation with DDI constraints
- Interaction-aware predictions
- Safety-focused drug selection

**SafeDrug** (`SafeDrug`)
- Safety-aware drug recommendation
- Molecular structure integration
- DDI constraint optimization
- Balances efficacy and safety

**Key Features:**
- Dual molecular graph encoders (global MPNN + local bipartite substructure encoder)
- DDI-controllable loss that penalizes interacting drug pairs
- Published at IJCAI 2021 (Yang et al., "SafeDrug: Dual Molecular Graph Encoders for Recommending Effective and Safe Drug Combinations")

**Usage:**
```python
from pyhealth.models import SafeDrug

# SafeDrug derives the DDI matrix and molecule set from the dataset's ATC
# drug codes (RDKit is a PyHealth dependency); it requires the label key "drugs".
model = SafeDrug(
    dataset=sample_dataset,   # from a DrugRecommendation* task
    embedding_dim=64,
    hidden_dim=64,
)
```

**MoleRec** (`MoleRec`)
- Molecular-level drug recommendations
- Sub-structure reasoning
- Fine-grained medication selection

### Disease Progression Models

**StageNet** (`StageNet`)
- Disease stage-aware prediction
- Learns clinical stages automatically
- Stage-adaptive feature extraction
- Effective for chronic disease monitoring

**Architecture:**
- Stage-aware LSTM
- Dynamic stage transitions
- Time-decay mechanism

**Usage:**
```python
from pyhealth.models import StageNet

# Use with a StageNet-format task such as MortalityPredictionStageNetMIMIC4
# (inputs "icd_codes" + "labs"); the schema, not the constructor, sets the inputs.
model = StageNet(
    dataset=sample_dataset,
    chunk_size=128,
)
```

**Best for:**
- ICU mortality prediction
- Chronic disease progression
- Time-varying risk assessment

**Deepr** (`Deepr`)
- Convolutional network over sequences of medical-record codes
- Medical concept embeddings with visit separators
- Published in IEEE Journal of Biomedical and Health Informatics (2017)

### Advanced Sequential Models

**Agent** (`Agent`)
- "Dr. Agent": clinical prediction via mimicked second opinions
- Two policy-gradient agents choose which parts of the patient history to attend to (dynamic skip connections)

**GRASP** (`GRASP`)
- Health-status representation learning that incorporates knowledge from similar patients (clustered patient graph); AAAI 2021

### Physiological Signal Models

**SparcNet** (`SparcNet`)
- 1D dense convolutional network from the expert-level EEG classification study of seizures and rhythmic/periodic patterns (Neurology 2023)
- Use for EEG event/abnormality classification tasks

**ContraWR** (`ContraWR`)
- Supervised encoder from the ContraWR sleep-EEG work (STFT + 2D CNN)
- Use for sleep staging and other spectrogram-style signal tasks

### Record Linkage

**MedLink** (`MedLink`)
- De-identified patient health record linkage (Wu et al., KDD 2023) — matches records of the same patient across sources; it is not concept/entity normalization

### Generative Models

- **Synthetic EHR**: `HALO`, `PromptEHR`, `MedGAN`, `CorGAN`, and `GPT`, used with the `EHRGeneration*` tasks
- **Images**: `GAN` and `VAE` generate or reconstruct small (32–128 px) images; `VAE` and `Graph_TorchvisionModel` still take the legacy `feature_keys`/`label_key`/`mode` arguments

### Social Determinants of Health

**SdohClassifier** (`SdohClassifier`)
- Sentence-level classification of social determinants of health from clinical text (MIMIC-III-derived SDoH annotations)

## Model Selection Guidelines

### By Task Type

**Binary Classification** (Mortality, Readmission)
- Start with: Logistic Regression (baseline)
- Standard: RNN, Transformer
- Interpretable: RETAIN, AdaCare
- Advanced: StageNet

**Multi-Label Classification** (Drug Recommendation)
- Standard: CNN, RNN
- Healthcare-specific: GAMENet, SafeDrug, MICRON, MoleRec
- Graph-based: GAT, GCN

**Regression** (Length of Stay)
- Start with: MLP (baseline)
- Sequential: RNN, TCN
- Advanced: Transformer

**Multi-Class Classification** (Medical Coding, Specialty)
- Standard: CNN, RNN, Transformer
- Text-based: TransformersModel (BERT variants)

### By Data Type

**Sequential Events** (Diagnoses, Medications, Procedures)
- RNN, LSTM, GRU
- Transformer
- RETAIN, AdaCare, ConCare

**Time-Series Signals** (EEG, ECG)
- CNN, TCN
- RNN
- Transformer

**Text** (Clinical Notes)
- TransformersModel (ClinicalBERT, BioBERT)
- CNN for shorter text
- RNN for sequential text

**Graphs** (Drug Interactions, Patient Networks)
- GNN (GAT, GCN)
- GAMENet, SafeDrug

**Images** (X-rays, CT scans)
- CNN (ResNet, DenseNet via TransformersModel)
- Vision Transformers

### By Interpretability Needs

**High Interpretability Required:**
- Logistic Regression
- RETAIN
- AdaCare
- SparcNet

**Moderate Interpretability:**
- CNN (filter visualization)
- Transformer (attention visualization)
- GNN (graph attention)

**Black-Box Acceptable:**
- Deep RNN models
- Complex ensembles

## Training Considerations

### Hyperparameter Tuning

**Embedding Dimension:**
- Small datasets: 64-128
- Large datasets: 128-256
- Complex tasks: 256-512

**Hidden Dimension:**
- Proportional to embedding_dim
- Typically 1-2x embedding_dim

**Number of Layers:**
- Start with 2-3 layers
- Deeper for complex patterns
- Watch for overfitting

**Dropout:**
- Start with 0.5
- Reduce if underfitting (0.1-0.3)
- Increase if overfitting (0.5-0.7)

### Computational Requirements

**Memory (GPU):**
- CNN: Low to moderate
- RNN: Moderate (sequence length dependent)
- Transformer: High (quadratic in sequence length)
- GNN: Moderate to high (graph size dependent)

**Training Speed:**
- Fastest: Logistic Regression, MLP, CNN
- Moderate: RNN, GNN
- Slower: Transformer (but parallelizable)

### Best Practices

1. **Start with simple baselines** (Logistic Regression, MLP)
2. **Choose the task's input schema** (the model's inputs) based on data availability
3. **Set the task's output schema to the prediction target** (binary, multiclass, multilabel, regression); the model's mode follows from it
4. **Consider interpretability requirements** for clinical deployment
5. **Validate on held-out test set** for realistic performance
6. **Monitor for overfitting** especially with complex models
7. **Use pretrained models** when possible (TransformersModel)
8. **Consider computational constraints** for deployment

## Example Workflow

```python
from pyhealth.datasets import MIMIC4EHRDataset
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer

# 1. Prepare data
dataset = MIMIC4EHRDataset(
    root="/path/to/data",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())

# 2. Initialize model
model = Transformer(
    dataset=sample_dataset,
    embedding_dim=128,
    num_layers=2,
    dropout=0.3,
)

# 3. Train model
trainer = Trainer(model=model, metrics=["pr_auc", "roc_auc", "f1"])
trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    monitor="pr_auc",
    monitor_criterion="max",
)

# 4. Evaluate
results = trainer.evaluate(test_loader)
print(results)
```

preprocessing.md 15 KB

# PyHealth Data Preprocessing and Processors

## Overview

PyHealth provides comprehensive data processing utilities to transform raw healthcare data into model-ready formats. Processors handle feature extraction, sequence processing, signal transformation, and label preparation.

> **How processors are actually wired in PyHealth 2.x (read first).** You rarely import and compose processor objects by hand. Instead, a **task's `input_schema` / `output_schema` maps each key to a processor by string name**, and `set_task` builds the processors for you. Examples seen in the codebase: `"sequence"`, `"timeseries"`, `"stagenet"`, `"stagenet_tensor"`, `"binary"`, `"multiclass"`, `"multilabel"`, `"regression"`. The tuple form passes kwargs: `("stagenet", {"padding": 0})`.
>
> ```python
> class MyTask(BaseTask):
>     input_schema  = {"conditions": "sequence", "labs": "timeseries"}
>     output_schema = {"mortality": "binary"}
> ```
>
> The class-by-class catalog below is a **conceptual reference** to the kinds of processing PyHealth performs. Treat the exact class names, import paths, and constructor kwargs as **illustrative, not verified** for 2.0.x — confirm against `pyhealth.processors` / the installed source before importing a specific class. Prefer the schema-string approach above.

## Processor Base Class

All processors inherit from `Processor` with standard interface:

**Key Methods:**
- `__call__()`: Transform input data
- `get_input_info()`: Return processed input schema
- `get_output_info()`: Return processed output schema

## Core Processor Types

### Feature Processors

**FeatureProcessor** (`FeatureProcessor`)
- Base class for feature extraction
- Handles vocabulary building
- Embedding preparation
- Feature encoding

**Common Operations:**
- Medical code tokenization
- Categorical encoding
- Feature normalization
- Missing value handling

**Usage:**
```python
from pyhealth.data import FeatureProcessor

processor = FeatureProcessor(
    vocabulary="diagnoses",
    min_freq=5,  # Minimum code frequency
    max_vocab_size=10000
)

processed_features = processor(raw_features)
```

### Sequence Processors

**SequenceProcessor** (`SequenceProcessor`)
- Processes sequential clinical events
- Temporal ordering preservation
- Sequence padding/truncation
- Time gap encoding

**Key Features:**
- Variable-length sequence handling
- Temporal feature extraction
- Sequence statistics computation

**Parameters:**
- `max_seq_length`: Maximum sequence length (truncate if longer)
- `padding`: Padding strategy ("pre" or "post")
- `truncating`: Truncation strategy ("pre" or "post")

**Usage:**
```python
from pyhealth.data import SequenceProcessor

processor = SequenceProcessor(
    max_seq_length=100,
    padding="post",
    truncating="post"
)

# Process diagnosis sequences
processed_seq = processor(diagnosis_sequences)
```

**NestedSequenceProcessor** (`NestedSequenceProcessor`)
- Handles hierarchical sequences (e.g., visits containing events)
- Two-level processing (visit-level and event-level)
- Preserves nested structure

**Use Cases:**
- EHR with visits containing multiple events
- Multi-level temporal modeling
- Hierarchical attention models

**Structure:**
```python
# Input: [[visit1_events], [visit2_events], ...]
# Output: Processed nested sequences with proper padding
```

### Numeric Data Processors

**NestedFloatsProcessor** (`NestedFloatsProcessor`)
- Processes nested numeric arrays
- Lab values, vital signs, measurements
- Multi-level numeric features

**Operations:**
- Normalization
- Standardization
- Missing value imputation
- Outlier handling

**Usage:**
```python
from pyhealth.data import NestedFloatsProcessor

processor = NestedFloatsProcessor(
    normalization="z-score",  # or "min-max"
    fill_missing="mean"  # imputation strategy
)

processed_labs = processor(lab_values)
```

**TensorProcessor** (`TensorProcessor`)
- Converts data to PyTorch tensors
- Type handling (long, float, etc.)
- Device placement (CPU/GPU)

**Parameters:**
- `dtype`: Tensor data type
- `device`: Computation device

### Time-Series Processors

**TimeseriesProcessor** (`TimeseriesProcessor`)
- Handles temporal data with timestamps
- Time gap computation
- Temporal feature engineering
- Irregular sampling handling

**Extracted Features:**
- Time since previous event
- Time to next event
- Event frequency
- Temporal patterns

**Usage:**
```python
from pyhealth.data import TimeseriesProcessor

processor = TimeseriesProcessor(
    time_unit="hour",  # "day", "hour", "minute"
    compute_gaps=True,
    compute_frequency=True
)

processed_ts = processor(timestamps, events)
```

**SignalProcessor** (`SignalProcessor`)
- Physiological signal processing
- EEG, ECG, PPG signals
- Filtering and preprocessing

**Operations:**
- Bandpass filtering
- Artifact removal
- Segmentation
- Feature extraction (frequency, amplitude)

**Usage:**
```python
from pyhealth.data import SignalProcessor

processor = SignalProcessor(
    sampling_rate=256,  # Hz
    bandpass_filter=(0.5, 50),  # Hz range
    segment_length=30  # seconds
)

processed_signal = processor(raw_eeg_signal)
```

### Image Processors

**ImageProcessor** (`ImageProcessor`)
- Medical image preprocessing
- Normalization and resizing
- Augmentation support
- Format standardization

**Operations:**
- Resize to standard dimensions
- Normalization (mean/std)
- Windowing (for CT/MRI)
- Data augmentation

**Usage:**
```python
from pyhealth.data import ImageProcessor

processor = ImageProcessor(
    image_size=(224, 224),
    normalization="imagenet",  # or custom mean/std
    augmentation=True
)

processed_image = processor(raw_image)
```

## Label Processors

### Binary Classification

**BinaryLabelProcessor** (`BinaryLabelProcessor`)
- Binary classification labels (0/1)
- Handles positive/negative classes
- Class weighting for imbalance

**Usage:**
```python
from pyhealth.data import BinaryLabelProcessor

processor = BinaryLabelProcessor(
    positive_class=1,
    class_weight="balanced"
)

processed_labels = processor(raw_labels)
```

### Multi-Class Classification

**MultiClassLabelProcessor** (`MultiClassLabelProcessor`)
- Multi-class classification (mutually exclusive classes)
- Label encoding
- Class balancing

**Parameters:**
- `num_classes`: Number of classes
- `class_weight`: Weighting strategy

**Usage:**
```python
from pyhealth.data import MultiClassLabelProcessor

processor = MultiClassLabelProcessor(
    num_classes=5,  # e.g., sleep stages: W, N1, N2, N3, REM
    class_weight="balanced"
)

processed_labels = processor(raw_labels)
```

### Multi-Label Classification

**MultiLabelProcessor** (`MultiLabelProcessor`)
- Multi-label classification (multiple labels per sample)
- Binary encoding for each label
- Label co-occurrence handling

**Use Cases:**
- Drug recommendation (multiple drugs)
- ICD coding (multiple diagnoses)
- Comorbidity prediction

**Usage:**
```python
from pyhealth.data import MultiLabelProcessor

processor = MultiLabelProcessor(
    num_labels=100,  # total possible labels
    threshold=0.5  # prediction threshold
)

processed_labels = processor(raw_label_sets)
```

### Regression

**RegressionLabelProcessor** (`RegressionLabelProcessor`)
- Continuous value prediction
- Target scaling and normalization
- Outlier handling

**Use Cases:**
- Length of stay prediction
- Lab value prediction
- Risk score estimation

**Usage:**
```python
from pyhealth.data import RegressionLabelProcessor

processor = RegressionLabelProcessor(
    normalization="z-score",  # or "min-max"
    clip_outliers=True,
    outlier_std=3  # clip at 3 standard deviations
)

processed_targets = processor(raw_values)
```

## Specialized Processors

### Text Processing

**TextProcessor** (`TextProcessor`)
- Clinical text preprocessing
- Tokenization
- Vocabulary building
- Sequence encoding

**Operations:**
- Lowercasing
- Punctuation removal
- Medical abbreviation handling
- Token frequency filtering

**Usage:**
```python
from pyhealth.data import TextProcessor

processor = TextProcessor(
    tokenizer="word",  # or "sentencepiece", "bpe"
    lowercase=True,
    max_vocab_size=50000,
    min_freq=5
)

processed_text = processor(clinical_notes)
```

### Model-Specific Processors

**StageNetProcessor** (`StageNetProcessor`)
- Specialized preprocessing for StageNet model
- Chunk-based sequence processing
- Stage-aware feature extraction

**Usage:**
```python
from pyhealth.data import StageNetProcessor

processor = StageNetProcessor(
    chunk_size=128,
    num_stages=3
)

processed_data = processor(sequential_data)
```

**StageNetTensorProcessor** (`StageNetTensorProcessor`)
- Tensor conversion for StageNet
- Proper batching and padding
- Stage mask generation

### Raw Data Processing

**RawProcessor** (`RawProcessor`)
- Minimal preprocessing
- Pass-through for pre-processed data
- Custom preprocessing scenarios

**Usage:**
```python
from pyhealth.data import RawProcessor

processor = RawProcessor()
processed_data = processor(data)  # Minimal transformation
```

## Sample-Level Processing

**SampleProcessor** (`SampleProcessor`)
- Processes complete samples (input + output)
- Coordinates multiple processors
- End-to-end preprocessing pipeline

**Workflow:**
1. Apply input processors to features
2. Apply output processors to labels
3. Combine into model-ready samples

**Usage:**
```python
from pyhealth.data import SampleProcessor

processor = SampleProcessor(
    input_processors={
        "diagnoses": SequenceProcessor(max_seq_length=50),
        "medications": SequenceProcessor(max_seq_length=30),
        "labs": NestedFloatsProcessor(normalization="z-score")
    },
    output_processor=BinaryLabelProcessor()
)

processed_sample = processor(raw_sample)
```

## Dataset-Level Processing

**DatasetProcessor** (`DatasetProcessor`)
- Processes entire datasets
- Batch processing
- Parallel processing support
- Caching for efficiency

**Operations:**
- Apply processors to all samples
- Generate vocabulary from dataset
- Compute dataset statistics
- Save processed data

**Usage:**
```python
from pyhealth.data import DatasetProcessor

processor = DatasetProcessor(
    sample_processor=sample_processor,
    num_workers=4,  # parallel processing
    cache_dir="/path/to/cache"
)

processed_dataset = processor(raw_dataset)
```

## Common Preprocessing Workflows

### Workflow 1: EHR Mortality Prediction

```python
from pyhealth.data import (
    SequenceProcessor,
    BinaryLabelProcessor,
    SampleProcessor
)

# Define processors
input_processors = {
    "diagnoses": SequenceProcessor(max_seq_length=50),
    "medications": SequenceProcessor(max_seq_length=30),
    "procedures": SequenceProcessor(max_seq_length=20)
}

output_processor = BinaryLabelProcessor(class_weight="balanced")

# Combine into sample processor
sample_processor = SampleProcessor(
    input_processors=input_processors,
    output_processor=output_processor
)

# Process dataset
processed_samples = [sample_processor(s) for s in raw_samples]
```

### Workflow 2: Sleep Staging from EEG

```python
from pyhealth.data import (
    SignalProcessor,
    MultiClassLabelProcessor,
    SampleProcessor
)

# Signal preprocessing
signal_processor = SignalProcessor(
    sampling_rate=100,
    bandpass_filter=(0.3, 35),  # EEG frequency range
    segment_length=30  # 30-second epochs
)

# Label processing
label_processor = MultiClassLabelProcessor(
    num_classes=5,  # W, N1, N2, N3, REM
    class_weight="balanced"
)

# Combine
sample_processor = SampleProcessor(
    input_processors={"signal": signal_processor},
    output_processor=label_processor
)
```

### Workflow 3: Drug Recommendation

```python
from pyhealth.data import (
    SequenceProcessor,
    MultiLabelProcessor,
    SampleProcessor
)

# Input processing
input_processors = {
    "diagnoses": SequenceProcessor(max_seq_length=50),
    "previous_medications": SequenceProcessor(max_seq_length=40)
}

# Multi-label output (multiple drugs)
output_processor = MultiLabelProcessor(
    num_labels=150,  # number of possible drugs
    threshold=0.5
)

sample_processor = SampleProcessor(
    input_processors=input_processors,
    output_processor=output_processor
)
```

### Workflow 4: Length of Stay Prediction

```python
from pyhealth.data import (
    SequenceProcessor,
    NestedFloatsProcessor,
    RegressionLabelProcessor,
    SampleProcessor
)

# Process different feature types
input_processors = {
    "diagnoses": SequenceProcessor(max_seq_length=30),
    "procedures": SequenceProcessor(max_seq_length=20),
    "labs": NestedFloatsProcessor(
        normalization="z-score",
        fill_missing="mean"
    )
}

# Regression target
output_processor = RegressionLabelProcessor(
    normalization="log",  # log-transform LOS
    clip_outliers=True
)

sample_processor = SampleProcessor(
    input_processors=input_processors,
    output_processor=output_processor
)
```

## Best Practices

### Sequence Processing

1. **Choose appropriate max_seq_length**: Balance between context and computation
   - Short sequences (20-50): Fast, less context
   - Medium sequences (50-100): Good balance
   - Long sequences (100+): More context, slower

2. **Truncation strategy**:
   - "post": Keep most recent events (recommended for clinical prediction)
   - "pre": Keep earliest events

3. **Padding strategy**:
   - "post": Pad at end (standard)
   - "pre": Pad at beginning

### Feature Encoding

1. **Vocabulary size**: Limit to frequent codes
   - `min_freq=5`: Include codes appearing ≥5 times
   - `max_vocab_size=10000`: Cap total vocabulary size

2. **Handle rare codes**: Group into "unknown" category

3. **Missing values**:
   - Imputation (mean, median, forward-fill)
   - Indicator variables
   - Special tokens

### Normalization

1. **Numeric features**: Always normalize
   - Z-score: Standard scaling (mean=0, std=1)
   - Min-max: Range scaling [0, 1]

2. **Compute statistics on training set only**: Prevent data leakage

3. **Apply same normalization to val/test sets**

### Class Imbalance

1. **Use class weighting**: `class_weight="balanced"`

2. **Consider oversampling**: For very rare positive cases

3. **Evaluate with appropriate metrics**: AUROC, AUPRC, F1

### Performance Optimization

1. **Cache processed data**: Save preprocessing results

2. **Parallel processing**: Use `num_workers` for DataLoader

3. **Batch processing**: Process multiple samples at once

4. **Feature selection**: Remove low-information features

### Validation

1. **Check processed shapes**: Ensure correct dimensions

2. **Verify value ranges**: After normalization

3. **Inspect samples**: Manually review processed data

4. **Monitor memory usage**: Especially for large datasets

## Troubleshooting

### Common Issues

**Memory Error:**
- Reduce `max_seq_length`
- Use smaller batches
- Process data in chunks
- Enable caching to disk

**Slow Processing:**
- Enable parallel processing (`num_workers`)
- Cache preprocessed data
- Reduce feature dimensionality
- Use more efficient data types

**Shape Mismatch:**
- Check sequence lengths
- Verify padding configuration
- Ensure consistent processor settings

**NaN Values:**
- Handle missing data explicitly
- Check normalization parameters
- Verify imputation strategy

**Class Imbalance:**
- Use class weighting
- Consider oversampling
- Adjust decision threshold
- Use appropriate evaluation metrics

tasks.md 13.7 KB

# PyHealth Clinical Prediction Tasks

## Overview

PyHealth provides 20+ predefined clinical prediction tasks for common healthcare AI applications. Each task transforms raw patient data into structured input-output pairs for model training.

> **Naming convention (PyHealth 2.x).** Tasks are **classes** (mostly `{Task}{Dataset}`, e.g. `MortalityPredictionMIMIC4`, `ReadmissionPredictionMIMIC3`, `DrugRecommendationMIMIC3`). You **instantiate** the class and pass the instance to `set_task`. The names below are the class names exported by `pyhealth.tasks` in 2.0.2; the naming is not fully regular (e.g. `LengthOfStayPredictioneICU`, `MIMIC3ICD9Coding`, `EEGAbnormalTUAB`), so copy them rather than deriving them. A few legacy snake-case functions (`sleep_staging_isruc_fn`, `sleep_staging_shhs_fn`, `patient_linkage_mimic3_fn`, `drug_recommendation_*_fn`) are still exported for older signal/linkage workflows. When unsure, run `dir(pyhealth.tasks)` in the installed version.

## Task Structure

Each task subclasses `BaseTask` (`from pyhealth.tasks.base_task import BaseTask`) and defines:

- **task_name**: String identifier
- **input_schema**: Dict mapping feature keys to processor types (e.g. `{"conditions": "sequence"}`; tuple form `("stagenet", {"padding": 0})` passes processor kwargs)
- **output_schema**: Dict mapping the label key to its type (e.g. `{"mortality": "binary"}`)
- **`__call__(patient)`**: Returns the list of sample dicts for one patient

**Usage Pattern:**
```python
from pyhealth.datasets import MIMIC4EHRDataset
from pyhealth.tasks import MortalityPredictionMIMIC4

dataset = MIMIC4EHRDataset(
    root="/path/to/data",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())
```

## Electronic Health Record (EHR) Tasks

### Mortality Prediction

**Purpose:** Predict patient death risk at next visit or within specified timeframe

**MIMIC-III Mortality** (`MortalityPredictionMIMIC3`; multimodal variant `MultimodalMortalityPredictionMIMIC3`)
- Predicts death at next hospital visit
- Binary classification task
- Input: Historical diagnoses, procedures, medications
- Output: Binary label (deceased/alive)

**MIMIC-IV Mortality** (`MortalityPredictionMIMIC4`; multimodal variant `MultimodalMortalityPredictionMIMIC4`)
- Updated version for MIMIC-IV dataset
- Enhanced feature set
- Improved label quality

**eICU Mortality** (`MortalityPredictionEICU`, `MortalityPredictionEICU2`)
- Multi-center ICU mortality prediction
- Accounts for hospital-level variation

**OMOP Mortality** (`MortalityPredictionOMOP`)
- Standardized mortality prediction
- Works with OMOP common data model

**In-Hospital Mortality** (`InHospitalMortalityMIMIC4`; MEDS-format data: `InHospitalMortalityMEDS`)
- Predicts death during current hospitalization
- Real-time risk assessment
- Earlier prediction window than next-visit mortality

**StageNet Mortality** (`MortalityPredictionStageNetMIMIC4`; inputs `icd_codes` + `labs`)
- Specialized for StageNet model architecture
- Temporal stage-aware prediction

### Hospital Readmission Prediction

**Purpose:** Identify patients at risk of hospital readmission within specified timeframe (typically 30 days)

**MIMIC-III Readmission** (`ReadmissionPredictionMIMIC3`)
- 30-day readmission prediction
- Binary classification
- Input: Diagnosis history, medications, demographics
- Output: Binary label (readmitted/not readmitted)

**MIMIC-IV Readmission** (`ReadmissionPredictionMIMIC4`)
- Enhanced readmission features
- Improved temporal modeling

**eICU Readmission** (`ReadmissionPredictionEICU`)
- ICU-specific readmission risk
- Multi-site data

**OMOP Readmission** (`ReadmissionPredictionOMOP`)
- Standardized readmission prediction

### Length of Stay Prediction

**Purpose:** Estimate hospital stay duration for resource planning and patient management

**MIMIC-III Length of Stay** (`LengthOfStayPredictionMIMIC3`)
- Regression task
- Input: Admission diagnoses, vitals, demographics
- Output: Continuous value (days)

**MIMIC-IV Length of Stay** (`LengthOfStayPredictionMIMIC4`; StageNet format: `LengthOfStayStageNetMIMIC4`)
- Enhanced features for LOS prediction
- Better temporal granularity

**eICU Length of Stay** (`LengthOfStayPredictioneICU`)
- ICU stay duration prediction
- Multi-hospital data

**OMOP Length of Stay** (`LengthOfStayPredictionOMOP`)
- Standardized LOS prediction

### Drug Recommendation

**Purpose:** Suggest appropriate medications based on patient history and current conditions

**MIMIC-III Drug Recommendation** (`DrugRecommendationMIMIC3`)
- Multi-label classification
- Input: Diagnoses, previous medications, demographics
- Output: Set of recommended drug codes
- Considers drug-drug interactions

**MIMIC-IV Drug Recommendation** (`DrugRecommendationMIMIC4`)
- Updated medication data
- Enhanced interaction modeling

**eICU Drug Recommendation** (`DrugRecommendationEICU`)
- Critical care medication recommendations

**OMOP Drug Recommendation** (`DrugRecommendationOMOP`)
- Standardized drug recommendation

**Key Considerations:**
- Handles polypharmacy scenarios
- Multi-label prediction (multiple drugs per patient)
- Can integrate with SafeDrug/GAMENet models for safety-aware recommendations

## Specialized Clinical Tasks

### Medical Coding

**MIMIC-III ICD-9 Coding** (`MIMIC3ICD9Coding`)
- Assigns ICD-9 diagnosis/procedure codes to clinical notes
- Multi-label text classification
- Input: Clinical text/documentation
- Output: Set of ICD-9 codes
- Supports both diagnosis and procedure coding

### Patient Linkage

**MIMIC-III Patient Linking** (`PatientLinkageMIMIC3Task`; legacy function `patient_linkage_mimic3_fn`)
- Record matching and deduplication
- Binary classification (same patient or not)
- Input: Demographic and clinical features from two records
- Output: Match probability

## Physiological Signal Tasks

### Sleep Staging

**Purpose:** Classify sleep stages from EEG/physiological signals for sleep disorder diagnosis

**ISRUC Sleep Staging** (legacy function `sleep_staging_isruc_fn`)
- Multi-class classification (Wake, N1, N2, N3, REM)
- Input: Multi-channel EEG signals
- Output: Sleep stage per epoch (typically 30 seconds)

**SleepEDF Sleep Staging** (`SleepStagingSleepEDF`; legacy function `sleep_staging_sleepedf_fn`)
- Standard sleep staging task
- PSG signal processing

**SHHS Sleep Staging** (legacy function `sleep_staging_shhs_fn`)
- Large-scale sleep study data
- Population-level sleep analysis

**Standardized Labels:**
- Wake (W)
- Non-REM Stage 1 (N1)
- Non-REM Stage 2 (N2)
- Non-REM Stage 3 (N3/Deep Sleep)
- REM (Rapid Eye Movement)

### EEG Analysis

**Abnormality Detection** (`EEGAbnormalTUAB`)
- Binary classification (normal/abnormal EEG)
- Clinical screening application
- Input: Multi-channel EEG recordings
- Output: Binary label

**Event Detection** (`EEGEventsTUEV`)
- Identify specific EEG events (spikes, seizures)
- Multi-class classification
- Input: EEG time series
- Output: Event type and timing

**Seizure Detection** (no built-in TUSZ task class in 2.0.2)
- Write a custom `BaseTask` over your EEG loader, or use the TUEV event task
- Input: Continuous EEG
- Output: Seizure/non-seizure classification

## Medical Imaging Tasks

### COVID-19 Chest X-ray Classification

**COVID-19 CXR** (`COVID19CXRClassification`; also `ChestXray14BinaryClassification`, `ChestXray14MultilabelClassification`)
- Multi-class image classification
- Classes: COVID-19, bacterial pneumonia, viral pneumonia, normal
- Input: Chest X-ray images
- Output: Disease classification

## Text-Based Tasks

### Medical Transcription Classification

**Medical Specialty Classification** (`MedicalTranscriptionsClassification`)
- Classify clinical notes by medical specialty
- Multi-class text classification
- Input: Clinical transcription text
- Output: Medical specialty (Cardiology, Neurology, etc.)

## Custom Task Creation

### Creating Custom Tasks

Subclass `BaseTask`, declare `input_schema` / `output_schema`, and implement `__call__` to emit one flat sample dict per prediction. Sample keys must match the schema keys; PyHealth wires up the matching processors automatically.

```python
from datetime import datetime, timedelta
from typing import Any, Dict, List
from pyhealth.tasks import BaseTask


class ThirtyDayReadmissionMIMIC4(BaseTask):
    task_name: str = "ThirtyDayReadmissionMIMIC4"
    # The model reads its inputs and label from these schemas.
    input_schema: Dict[str, str] = {
        "conditions": "sequence",
        "procedures": "sequence",
    }
    output_schema: Dict[str, str] = {"readmitted": "binary"}

    def __call__(self, patient: Any) -> List[Dict[str, Any]]:
        samples: List[Dict[str, Any]] = []
        admissions = patient.get_events(event_type="admissions")

        for i in range(len(admissions) - 1):
            adm, nxt = admissions[i], admissions[i + 1]
            discharge = datetime.strptime(adm.dischtime, "%Y-%m-%d %H:%M:%S")

            # Events are table rows; filter them to this admission's time window
            diagnoses = patient.get_events(
                event_type="diagnoses_icd", start=adm.timestamp, end=discharge
            )
            procedures = patient.get_events(
                event_type="procedures_icd", start=adm.timestamp, end=discharge
            )
            conditions = [e.icd_code for e in diagnoses if getattr(e, "icd_code", None)]
            procs = [e.icd_code for e in procedures if getattr(e, "icd_code", None)]
            if not conditions or not procs:
                continue

            samples.append({
                "patient_id": patient.patient_id,
                "visit_id": adm.hadm_id,
                "conditions": conditions,
                "procedures": procs,
                "readmitted": int(nxt.timestamp - discharge <= timedelta(days=30)),
            })

        return samples


# Apply the custom task (instantiate it); the dataset must load the
# "diagnoses_icd" and "procedures_icd" tables
sample_dataset = dataset.set_task(ThirtyDayReadmissionMIMIC4())
```

> The pattern mirrors the built-in `MortalityPredictionMIMIC4`: `patient.get_events(event_type=<table>, start=..., end=...)` returns rows whose columns are attributes (`icd_code`, `hadm_id`, `dischtime`, ...). Column names come from the dataset's YAML config, so confirm them for other datasets before relying on a specific attribute.

### Task Function Components

1. **Input Schema Definition**
   - Specify which features to extract
   - Define feature types (codes, sequences, values)
   - Set temporal windows

2. **Output Schema Definition**
   - Define prediction targets
   - Set label types (binary, multi-class, multi-label, regression)
   - Specify evaluation metrics

3. **Filtering Logic**
   - Exclude patients/visits with insufficient data
   - Apply inclusion/exclusion criteria
   - Handle missing data

4. **Sample Generation**
   - Create input-output pairs
   - Maintain patient/visit identifiers
   - Preserve temporal ordering

## Task Selection Guidelines

### Clinical Prediction Tasks
**Use when:** Working with structured EHR data (diagnoses, medications, procedures)

**Datasets:** MIMIC-III, MIMIC-IV, eICU, OMOP

**Common tasks:**
- Mortality prediction for risk stratification
- Readmission prediction for care transition planning
- Length of stay for resource allocation
- Drug recommendation for clinical decision support

### Signal Processing Tasks
**Use when:** Working with physiological time-series data

**Datasets:** SleepEDF, SHHS, ISRUC, TUEV, TUAB

**Common tasks:**
- Sleep staging for sleep disorder diagnosis
- EEG abnormality detection for screening
- Seizure detection for epilepsy monitoring

### Imaging Tasks
**Use when:** Working with medical images

**Datasets:** COVID-19 CXR

**Common tasks:**
- Disease classification from radiographs
- Abnormality detection

### Text Tasks
**Use when:** Working with clinical notes and documentation

**Datasets:** Medical Transcriptions, MIMIC-III (with notes)

**Common tasks:**
- Medical coding from clinical text
- Specialty classification
- Clinical information extraction

## Task Output Structure

`set_task` returns a `SampleDataset` of flat sample dicts whose keys match the task's `input_schema` / `output_schema`:

```python
sample = {
    "patient_id": "unique_patient_id",
    # feature keys from input_schema, e.g.:
    "conditions": ["428.0", "401.9"],
    "procedures": ["9904", "3893"],
    "drugs": ["Metoprolol", "Lisinopril"],
    # label key from output_schema, e.g.:
    "mortality": 0,
}
```

## Integration with Models

Models read their inputs and label directly from the task's `input_schema` / `output_schema`, so no key arguments are passed:

```python
from pyhealth.datasets import MIMIC4EHRDataset
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer

# 1. Create task-specific dataset
dataset = MIMIC4EHRDataset(
    root="/path/to/data",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())

# 2. The model picks up "conditions"/"procedures"/"drugs" -> "mortality" (binary)
model = Transformer(dataset=sample_dataset, embedding_dim=128)
```

## Best Practices

1. **Match task to clinical question**: Choose predefined tasks when available for standardized benchmarking

2. **Consider temporal windows**: Ensure sufficient history for meaningful predictions

3. **Handle class imbalance**: Many clinical outcomes are rare (mortality, readmission)

4. **Validate clinical relevance**: Ensure prediction windows align with clinical decision-making timelines

5. **Use appropriate metrics**: Different tasks require different evaluation metrics (AUROC for binary, macro-F1 for multi-class)

6. **Document exclusion criteria**: Track which patients/visits are filtered and why

7. **Preserve patient privacy**: Always use de-identified data and follow HIPAA/GDPR guidelines

training_evaluation.md 20.4 KB

# PyHealth Training, Evaluation, and Interpretability

## Overview

PyHealth provides comprehensive tools for training models, evaluating predictions, ensuring model reliability, and interpreting results for clinical applications.

## Trainer Class

### Core Functionality

The `Trainer` class manages the complete model training and evaluation workflow with PyTorch integration.

**Initialization:**
```python
from pyhealth.trainer import Trainer

trainer = Trainer(
    model=model,                          # PyHealth model
    metrics=["pr_auc", "roc_auc", "f1"],  # metrics computed by evaluate()
    # device is auto-detected; pass device="cpu"/"cuda" to override
)
```

The metric list you pass here is what `evaluate()` reports and what `monitor=` can reference during training.

### Training

**train() method**

Trains models with comprehensive monitoring and checkpointing.

**Parameters:**
- `train_dataloader`: Training data loader
- `val_dataloader`: Validation data loader (optional)
- `epochs`: Number of training epochs
- `optimizer_class`: Optimizer **class** (e.g. `torch.optim.Adam`, `torch.optim.AdamW`)
- `optimizer_params`: Dict of optimizer kwargs (e.g. `{"lr": 1e-3, "weight_decay": 1e-5}`)
- `monitor`: Metric to monitor — one of the names passed to `Trainer(metrics=...)`, e.g. `"pr_auc"`
- `monitor_criterion`: "max" or "min"

**Usage:**
```python
import torch

trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    optimizer_class=torch.optim.Adam,
    optimizer_params={"lr": 1e-3, "weight_decay": 1e-5},
    monitor="pr_auc",
    monitor_criterion="max",
)
```

**Training Features:**

1. **Automatic Checkpointing**: Saves the best model by the monitored metric (`best.ckpt`, when logging is enabled)

2. **Early Stopping**: `patience=<epochs>` stops training when the monitored metric stops improving

3. **Gradient Clipping**: `max_grad_norm=<float>` clips gradients

4. **Progress Tracking**: Logs training progress and validation metrics

5. **Device Placement**: Single device — CUDA if available, else CPU (override with `Trainer(device=...)`); multi-GPU training is not built in

### Inference

**inference() method**

Performs predictions on datasets.

**Parameters:**
- `dataloader`: Data loader for inference
- `additional_outputs`: List of additional outputs to return
- `return_patient_ids`: Return patient identifiers

**Usage:**
```python
# Default: returns a 3-tuple
y_true, y_prob, loss = trainer.inference(test_loader)

# With patient IDs: returns a 4-tuple
y_true, y_prob, loss, patient_ids = trainer.inference(test_loader, return_patient_ids=True)

# additional_outputs=[...] inserts a dict before patient_ids; request only keys the
# model's forward() actually returns (e.g. "logit"; calibrated set models add "y_predset")
y_true, y_prob, loss, extra, patient_ids = trainer.inference(
    test_loader, additional_outputs=["logit"], return_patient_ids=True,
)
```

**Returns (tuple, not a dict):**
- `y_true`: Ground truth labels (array)
- `y_prob`: Predicted probabilities (array)
- `loss` / `mean_loss`: Mean loss over the dataset
- `additional_outputs`: Dict of requested extras (only if `additional_outputs=` given)
- `patient_ids`: Patient identifiers (only if `return_patient_ids=True`)

### Evaluation

**evaluate() method**

Computes comprehensive evaluation metrics.

**Parameters:**
- `dataloader`: Data loader for evaluation

`evaluate()` uses the metric list set on the `Trainer` (`Trainer(metrics=[...])`) — it does **not** take a `metrics=` argument. To compute metrics ad hoc, call the metric function directly on `inference()` outputs.

**Usage:**
```python
# Uses metrics passed to Trainer(...)
results = trainer.evaluate(test_loader)
print(results)
# e.g. {'pr_auc': 0.78, 'roc_auc': 0.82, 'f1': 0.73}

# Or compute metrics manually from predictions
from pyhealth.metrics.binary import binary_metrics_fn

y_true, y_prob, loss = trainer.inference(test_loader)
binary_metrics_fn(y_true, y_prob, metrics=["pr_auc", "roc_auc", "f1"])
```

### Checkpoint Management

The `Trainer` saves and restores the model's `state_dict` (there is no `trainer.save()` / `trainer.load()`):

```python
trainer.save_ckpt("./models/best_model.pt")
trainer.load_ckpt("./models/best_model.pt")

# Or load while constructing a trainer for an identically configured model
trainer = Trainer(model=model, checkpoint_path="./models/best_model.pt")
```

With logging enabled (the default), `train()` writes `last.ckpt` and `best.ckpt` (by `monitor`) under `output_path/exp_name` and reloads `best.ckpt` at the end (`load_best_model_at_last=True`). With `enable_logging=False` nothing is written, so call `save_ckpt()` yourself.

## Evaluation Metrics

### Binary Classification Metrics

**Available metric strings:**
- `accuracy`: Overall accuracy
- `f1`: F1 score
- `precision`, `recall`, `balanced_accuracy`, `jaccard`
- `roc_auc`: Area under ROC curve
- `pr_auc`: Area under precision-recall curve
- `cohen_kappa`: Inter-rater reliability

(Strings have no `_score` suffix.)

**Usage:**
```python
from pyhealth.metrics.binary import binary_metrics_fn

# Note: arg is y_prob (predicted probabilities), not thresholded y_pred
metrics = binary_metrics_fn(
    y_true=labels,
    y_prob=probabilities,
    metrics=["accuracy", "f1", "pr_auc", "roc_auc"],
)
```

**Threshold Selection:**
```python
# Default threshold: 0.5
predictions_binary = (predictions > 0.5).astype(int)

# Optimal threshold by F1
from sklearn.metrics import f1_score
thresholds = np.arange(0.1, 0.9, 0.05)
f1_scores = [f1_score(y_true, (y_pred > t).astype(int)) for t in thresholds]
optimal_threshold = thresholds[np.argmax(f1_scores)]
```

**Best Practices:**
- **Use AUROC**: Overall model discrimination
- **Use AUPRC**: Especially for imbalanced classes
- **Use F1**: Balance precision and recall
- **Report confidence intervals**: Bootstrap resampling

### Multi-Class Classification Metrics

**Available metric strings:**
- `accuracy`: Overall accuracy
- `f1_macro`: Unweighted mean F1 across classes
- `f1_micro`: Global F1 (total TP, FP, FN)
- `f1_weighted`: Weighted mean F1 by class frequency
- `cohen_kappa`: Multi-class kappa

**Usage:**
```python
from pyhealth.metrics.multiclass import multiclass_metrics_fn

metrics = multiclass_metrics_fn(
    y_true=labels,
    y_prob=probabilities,
    metrics=["accuracy", "f1_macro", "f1_weighted"],
)
```

**Per-Class Metrics:**
```python
from sklearn.metrics import classification_report

print(classification_report(y_true, y_pred,
    target_names=["Wake", "N1", "N2", "N3", "REM"]))
```

**Confusion Matrix:**
```python
from sklearn.metrics import confusion_matrix
import seaborn as sns

cm = confusion_matrix(y_true, y_pred)
sns.heatmap(cm, annot=True, fmt='d')
```

### Multi-Label Classification Metrics

**Available metric strings** (the `*_samples` family is the per-example average used for drug recommendation):
- `jaccard_samples`: Sample-averaged Jaccard (intersection over union)
- `f1_samples`: Sample-averaged F1
- `pr_auc_samples`: Sample-averaged AUPRC
- `hamming_loss`: Fraction of incorrect labels
- `ddi`: Drug-drug interaction rate (drug-rec models; returned under the key `ddi_score`)

**Usage:**
```python
from pyhealth.metrics.multilabel import multilabel_metrics_fn

# y_prob: [n_samples, n_labels] probability matrix
metrics = multilabel_metrics_fn(
    y_true=label_matrix,
    y_prob=prob_matrix,
    metrics=["jaccard_samples", "f1_samples", "pr_auc_samples"],
)
```

**Drug Recommendation Metrics:**
```python
# Jaccard similarity (intersection/union)
jaccard = len(set(true_drugs) & set(pred_drugs)) / len(set(true_drugs) | set(pred_drugs))

# Precision@k: Precision for top-k predictions
def precision_at_k(y_true, y_pred, k=10):
    top_k_pred = y_pred.argsort()[-k:]
    return len(set(y_true) & set(top_k_pred)) / k
```

### Regression Metrics

**Available metric strings (2.0.2):**
- `mae`: Mean absolute error
- `mse`: Mean squared error
- `kl_divergence`: KL divergence between the normalized target and prediction vectors

Any other name (e.g. `rmse`, `r2`) raises `ValueError`; compute those with scikit-learn.

**Usage:**
```python
from pyhealth.metrics.regression import regression_metrics_fn

# Signature: regression_metrics_fn(x, x_rec, metrics=None) — positional true, predicted
metrics = regression_metrics_fn(true_values, predictions, metrics=["mae", "mse"])

from sklearn.metrics import r2_score, root_mean_squared_error
extra = {"rmse": root_mean_squared_error(true_values, predictions), "r2": r2_score(true_values, predictions)}
```

**Percentage Error Metrics:**
```python
# Mean Absolute Percentage Error
mape = np.mean(np.abs((y_true - y_pred) / y_true)) * 100

# Median Absolute Percentage Error (robust to outliers)
medape = np.median(np.abs((y_true - y_pred) / y_true)) * 100
```

### Fairness Metrics

**Purpose:** Assess model bias across a protected vs. unprotected group.

**`pyhealth.metrics.fairness.fairness_metrics_fn`** — available metric strings:
- `disparate_impact`: Ratio of favorable-outcome rates (protected / unprotected)
- `statistical_parity_difference`: Difference in favorable-outcome rates

**Signature:** `fairness_metrics_fn(y_true, y_prob, sensitive_attributes, favorable_outcome=1, metrics=[...], threshold=0.5)` where `sensitive_attributes` is a 0/1 array (1 = protected group).

**Usage:**
```python
from pyhealth.metrics.fairness import fairness_metrics_fn

# sensitive_attributes: 1 for the protected group, 0 otherwise
fairness_results = fairness_metrics_fn(
    y_true=labels,
    y_prob=probabilities,
    sensitive_attributes=protected_mask,
    favorable_outcome=1,
    metrics=["disparate_impact", "statistical_parity_difference"],
)
```

**Example:**
```python
# Evaluate fairness across gender
male_mask = (demographics == "male")
female_mask = (demographics == "female")

male_tpr = recall_score(y_true[male_mask], y_pred[male_mask])
female_tpr = recall_score(y_true[female_mask], y_pred[female_mask])

tpr_disparity = abs(male_tpr - female_tpr)
print(f"TPR disparity: {tpr_disparity:.3f}")
```

## Calibration and Uncertainty Quantification

### Model Calibration

**Purpose:** Ensure predicted probabilities match actual frequencies

**Calibration Plot:**
```python
from sklearn.calibration import calibration_curve
import matplotlib.pyplot as plt

fraction_of_positives, mean_predicted_value = calibration_curve(
    y_true, y_prob, n_bins=10
)

plt.plot(mean_predicted_value, fraction_of_positives, marker='o')
plt.plot([0, 1], [0, 1], linestyle='--', label='Perfect calibration')
plt.xlabel('Mean predicted probability')
plt.ylabel('Fraction of positives')
plt.legend()
```

**Expected Calibration Error (ECE):**
```python
def expected_calibration_error(y_true, y_prob, n_bins=10):
    """Compute ECE"""
    bins = np.linspace(0, 1, n_bins + 1)
    bin_indices = np.digitize(y_prob, bins) - 1

    ece = 0
    for i in range(n_bins):
        mask = bin_indices == i
        if mask.sum() > 0:
            bin_accuracy = y_true[mask].mean()
            bin_confidence = y_prob[mask].mean()
            ece += mask.sum() / len(y_true) * abs(bin_accuracy - bin_confidence)

    return ece
```

**Calibration Methods:**

1. **Platt Scaling**: Logistic regression on validation predictions
```python
from sklearn.linear_model import LogisticRegression

calibrator = LogisticRegression()
calibrator.fit(val_predictions.reshape(-1, 1), val_labels)
calibrated_probs = calibrator.predict_proba(test_predictions.reshape(-1, 1))[:, 1]
```

2. **Isotonic Regression**: Non-parametric calibration
```python
from sklearn.isotonic import IsotonicRegression

calibrator = IsotonicRegression(out_of_bounds='clip')
calibrator.fit(val_predictions, val_labels)
calibrated_probs = calibrator.predict(test_predictions)
```

3. **Temperature Scaling and other built-in calibrators**: PyHealth wraps a trained model with a calibrator from `pyhealth.calib.calibration` (`TemperatureScaling`, `HistogramBinning`, `DirichletCalibration`, `KCal`) and fits it on a held-out calibration split. Check each class docstring for the task modes it supports.
```python
from pyhealth.calib.calibration import TemperatureScaling

cal_model = TemperatureScaling(model)       # model = trained PyHealth model
cal_model.calibrate(cal_dataset=val_data)   # a split not used for checkpoint selection, ideally
print(Trainer(model=cal_model, metrics=["accuracy"]).evaluate(test_loader))
```

### Uncertainty Quantification

**Conformal Prediction:**

Provide prediction sets with guaranteed coverage.

**Usage (split conformal, multiclass, plain NumPy):**
```python
import numpy as np

alpha = 0.1                                   # target 90% coverage
n = len(cal_labels)                           # held-out calibration split
scores = 1 - cal_probs[np.arange(n), cal_labels]
q_level = np.ceil((n + 1) * (1 - alpha)) / n  # finite-sample correction
qhat = np.quantile(scores, q_level, method="higher")

prediction_sets = test_probs >= (1 - qhat)    # boolean [n_test, n_classes]
coverage = prediction_sets[np.arange(len(test_labels)), test_labels].mean()
avg_size = prediction_sets.sum(axis=1).mean()
```

PyHealth also ships prediction-set constructors in `pyhealth.calib.predictionset` (`LABEL`, `SCRIB`, `FavMac`, `CovariateLabel`, `ClusterLabel`, `NeighborhoodLabel`). They wrap a trained model, are fit with `.calibrate(cal_dataset=...)`, and expose the sets via `Trainer(model=cal_model).inference(loader, additional_outputs=["y_predset"])`; the multiclass metric function accepts them through `y_predset=` with metrics such as `miscoverage_ps`, `set_size`, and `rejection_rate`.

**Monte Carlo Dropout:**

Estimate uncertainty through dropout at inference.

```python
def predict_with_uncertainty(model, dataloader, num_samples=20):
    """Predict with uncertainty using MC dropout"""
    model.train()  # Keep dropout active

    predictions = []
    for _ in range(num_samples):
        batch_preds = []
        for batch in dataloader:
            with torch.no_grad():
                output = model(**batch)["y_prob"]
                batch_preds.append(output)
        predictions.append(torch.cat(batch_preds))

    predictions = torch.stack(predictions)
    mean_pred = predictions.mean(dim=0)
    std_pred = predictions.std(dim=0)  # Uncertainty

    return mean_pred, std_pred
```

**Ensemble Uncertainty:**

```python
# Train multiple models
models = [train_model(seed=i) for i in range(5)]

# Predict with ensemble
ensemble_preds = []
for model in models:
    pred = model.predict(test_data)
    ensemble_preds.append(pred)

mean_pred = np.mean(ensemble_preds, axis=0)
std_pred = np.std(ensemble_preds, axis=0)  # Uncertainty
```

## Interpretability

### Attention and Attribution Methods

PyHealth models do not return raw attention matrices from `forward()` / `Trainer.inference()` (the output dict is `loss`, `y_prob`, `y_true`, `logit`). Use the interpreters in `pyhealth.interpret.methods` instead; each takes a trained model and returns a dict of per-token scores keyed by input feature (`attribute(**batch)`):

| Method | Works with (2.0.2) |
|--------|--------------------|
| `CheferRelevance`, `AttentionRollout` | Attention models: `Transformer`, `StageAttentionNet` |
| `IntegratedGradients`, `DeepLift` (`use_embeddings=True`) | Models with `forward_from_embedding`: `Transformer`, `MLP`, `StageNet`, `StageAttentionNet`, `TorchvisionModel` |
| `ShapExplainer`, `LimeExplainer` | Perturbation-based; see the class docstrings for input requirements |

```python
from pyhealth.interpret.methods import IntegratedGradients

ig = IntegratedGradients(model, use_embeddings=True, steps=50)
batch = next(iter(get_dataloader(test_data, batch_size=1, shuffle=False)))
attributions = ig.attribute(**batch)          # {"conditions": tensor, ...}
```

**RETAIN** is interpretable by design (visit-level alpha and variable-level beta attention), but PyHealth 2.0.2 does not expose those weights through `forward()`, and RETAIN implements neither the Chefer interface nor `forward_from_embedding`. For attributions you can hand to clinicians, train a `Transformer` alongside it.

### Feature Importance

**Permutation Importance:**

```python
from sklearn.inspection import permutation_importance

def get_predictions(model, X):
    return model.predict(X)

result = permutation_importance(
    model, X_test, y_test,
    n_repeats=10,
    scoring='roc_auc'
)

# Sort features by importance
indices = result.importances_mean.argsort()[::-1]
for i in indices[:10]:
    print(f"{feature_names[i]}: {result.importances_mean[i]:.3f}")
```

**SHAP Values:**

```python
import shap

# Create explainer
explainer = shap.DeepExplainer(model, train_data)

# Compute SHAP values
shap_values = explainer.shap_values(test_data)

# Visualize
shap.summary_plot(shap_values, test_data, feature_names=feature_names)
```

### Chefer Relevance (PyHealth's built-in attention interpretability)

PyHealth ships the Chefer relevance method for attention-based models (`Transformer`, `StageAttentionNet`; other models raise `ValueError`). The class is `CheferRelevance` in `pyhealth.interpret.methods`; call `attribute(**batch)` (the older `get_relevance_matrix(**batch)` is a deprecated alias).

```python
from pyhealth.interpret.methods import CheferRelevance
from pyhealth.datasets import get_dataloader

relevance = CheferRelevance(model)

# One sample at a time (batch_size=1)
loader = get_dataloader(test_dataset, batch_size=1, shuffle=False)
batch = next(iter(loader))

scores = relevance.attribute(**batch)  # dict: feature_key -> [batch, num_tokens] tensor
for feature_key, rel in scores.items():
    top_tokens = rel[0].topk(5).indices
    print(f"{feature_key}: top-5 tokens -> {top_tokens.tolist()}")
```

## Complete Training Pipeline Example

```python
import torch
from pyhealth.datasets import MIMIC4EHRDataset, split_by_patient, get_dataloader
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer

# 1. Load and prepare data
dataset = MIMIC4EHRDataset(
    root="/path/to/mimic4",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())

# 2. Split data by patient
train_data, val_data, test_data = split_by_patient(
    sample_dataset, [0.7, 0.1, 0.2]
)

# 3. Create data loaders
train_loader = get_dataloader(train_data, batch_size=64, shuffle=True)
val_loader = get_dataloader(val_data, batch_size=64, shuffle=False)
test_loader = get_dataloader(test_data, batch_size=64, shuffle=False)

# 4. Initialize model
model = Transformer(
    dataset=sample_dataset,
    embedding_dim=128,
    num_layers=2,
    dropout=0.3,
)

# 5. Train model
trainer = Trainer(model=model, metrics=["accuracy", "pr_auc", "roc_auc", "f1"])
trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    optimizer_class=torch.optim.Adam,
    optimizer_params={"lr": 1e-3, "weight_decay": 1e-5},
    monitor="pr_auc",
    monitor_criterion="max",
)

# 6. Evaluate on test set (uses the Trainer's metric list)
test_results = trainer.evaluate(test_loader)
print("Test Results:")
for metric, value in test_results.items():
    print(f"{metric}: {value:.4f}")

# 7. Get predictions for analysis (tuple unpacking)
y_true, y_prob, loss = trainer.inference(test_loader)

# 8. Calibration analysis
from sklearn.calibration import calibration_curve

positive_prob = y_prob if y_prob.ndim == 1 else y_prob[..., -1]
fraction_pos, mean_pred = calibration_curve(y_true, positive_prob, n_bins=10)
ece = expected_calibration_error(y_true, positive_prob)
print(f"Expected Calibration Error: {ece:.4f}")

# 9. Save final model
trainer.save_ckpt("./models/mortality_transformer_final.pt")
```

## Best Practices

### Training

1. **Monitor multiple metrics**: Track both loss and task-specific metrics
2. **Use validation set**: Prevent overfitting with early stopping
3. **Gradient clipping**: Stabilize training (max_grad_norm=5.0)
4. **Learning rate scheduling**: Reduce LR on plateau
5. **Checkpoint best model**: Save based on validation performance

### Evaluation

1. **Use task-appropriate metrics**: AUROC/AUPRC for binary, macro-F1 for imbalanced multi-class
2. **Report confidence intervals**: Bootstrap or cross-validation
3. **Stratified evaluation**: Report metrics by subgroups
4. **Clinical metrics**: Include clinically relevant thresholds
5. **Fairness assessment**: Evaluate across demographic groups

### Deployment

1. **Calibrate predictions**: Ensure probabilities are reliable
2. **Quantify uncertainty**: Provide confidence estimates
3. **Monitor performance**: Track metrics in production
4. **Handle distribution shift**: Detect when data changes
5. **Interpretability**: Provide explanations for predictions

SKILL.md 18.9 KB

---
name: alterlab-pyhealth
description: Develops, tests, and validates clinical machine learning models with the PyHealth 2.x healthcare AI toolkit. Use when working with electronic health records (EHR), clinical prediction tasks (mortality, readmission, length of stay, drug recommendation), medical coding systems (ICD, NDC, ATC, CCS), physiological signals (EEG, ECG), healthcare datasets (MIMIC-III/IV, eICU, OMOP), or implementing deep learning models for healthcare (RETAIN, SafeDrug, GAMENet, Transformer, GAT/GCN). Part of the AlterLab Academic Skills suite.
license: MIT
allowed-tools: Read Write Edit Bash(python:*)
compatibility: "Self-contained — runs under `uv run python` with PyHealth >= 2.0.2 (Python 3.12-3.13) in its own environment; no API key required. MIMIC-III/IV and eICU data need the user's own PhysioNet credentialed access."
metadata:
    skill-author: AlterLab
    version: "1.2.0"
    last_updated: "2026-09-23"
---

# PyHealth: Healthcare AI Toolkit

## Overview

PyHealth is a Python library for healthcare AI that provides datasets, task definitions, models, trainers, and medical-code utilities for clinical machine learning. Use this skill when developing healthcare prediction models, processing clinical data, working with medical coding systems, or validating models before any clinical use.

> **Version gotcha (read first).** This skill targets **PyHealth 2.x** (current release 2.0.2, Sept 2026). The 2.0 rewrite changed the API in ways most tutorials and pre-2025 snippets get wrong:
> - **Tasks are classes you instantiate**, e.g. `MortalityPredictionMIMIC4()`, `DrugRecommendationMIMIC3()` — not the old snake-case `mortality_prediction_mimic4_fn` functions. Pass the instance to `dataset.set_task(task)`.
> - **Datasets take an explicit table list.** Single-source loaders use `root=` + `tables=[...]` (`MIMIC3Dataset`, `MIMIC4EHRDataset`, `eICUDataset`, `OMOPDataset`); the multimodal `MIMIC4Dataset` uses `ehr_root=` + `ehr_tables=[...]` (plus optional `note_root`/`cxr_root`).
> - **Models take only the `SampleDataset` plus hyperparameters**, e.g. `Transformer(dataset=samples, embedding_dim=128)`. Feature keys, label key, and mode are read from the task's `input_schema` / `output_schema`; the 1.x `feature_keys=` / `label_key=` / `mode=` arguments raise `TypeError`.
> - **Metric names have no `_score` suffix**: `pr_auc`, `roc_auc`, `f1`; multilabel/drug-rec use the `*_samples` family (`jaccard_samples`, `f1_samples`, `pr_auc_samples`). Pass `metrics=[...]` to the **`Trainer` constructor** and `monitor=` one of those names.
> - Checkpoints use `trainer.save_ckpt(path)` / `trainer.load_ckpt(path)` (there is no `trainer.save`).
> - 2.0.2 requires **Python 3.12 or 3.13** and pins its own stack (numpy 2.2, pandas 2.3, torch 2.7, transformers 4.53), so install it in a dedicated environment rather than next to pandas 3 / transformers 5.
>
> When unsure of a class or argument name, check the installed source rather than trusting older snippets.

## When to Use This Skill

Invoke this skill when:

- **Working with healthcare datasets**: MIMIC-III, MIMIC-IV, eICU, OMOP, sleep EEG data, medical images
- **Clinical prediction tasks**: Mortality prediction, hospital readmission, length of stay, drug recommendation
- **Medical coding**: Translating between ICD-9/10, NDC, RxNorm, ATC, CCS coding systems
- **Processing clinical data**: Sequential events, physiological signals, clinical text, medical images
- **Implementing healthcare models**: RETAIN, SafeDrug, GAMENet, StageNet, Transformer for EHR
- **Evaluating clinical models**: Fairness metrics, calibration, interpretability, uncertainty quantification

### Does NOT Trigger

| Scenario | Use Instead |
|----------|-------------|
| Cleaning a raw ECG/EEG/EDA trace and computing HRV or SCR features (no model training) | `alterlab-neurokit2` |
| Reading, anonymizing, or converting DICOM image files | `alterlab-pydicom` |
| Kaplan-Meier / Cox time-to-event modeling on a tabular clinical dataset | `alterlab-scikit-survival` |
| Biomarker-stratified cohort report with GRADE-graded treatment recommendations | `alterlab-clinical-decision` |
| General tabular ML on non-EHR data (scikit-learn pipelines) | `alterlab-scikit-learn` |

## Core Capabilities

PyHealth operates through a modular 5-stage pipeline:

1. **Data Loading**: Standardized loaders for EHR, signal, imaging, and text datasets
2. **Task Definition**: Predefined clinical prediction tasks (task classes) or custom `BaseTask` subclasses
3. **Model Selection**: Baselines, general deep learning, and healthcare-specific models
4. **Training**: `Trainer` with best-checkpoint selection, monitoring, and evaluation
5. **Validation**: Calibration, conformal prediction, fairness metrics, and interpretability methods

PyHealth 2.x uses a **polars-backed** data layer and caches task samples, which keeps large EHR tables memory-efficient.

## Quick Start Workflow

```python
from pyhealth.datasets import MIMIC4EHRDataset, split_by_patient, get_dataloader
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer

# 1. Load dataset (declare the tables the task needs) and set the task (a class instance)
dataset = MIMIC4EHRDataset(
    root="/path/to/mimic-iv/2.2",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())

# 2. Split data by patient (no leakage across splits)
train, val, test = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=42)

# 3. Create data loaders
train_loader = get_dataloader(train, batch_size=64, shuffle=True)
val_loader = get_dataloader(val, batch_size=64, shuffle=False)
test_loader = get_dataloader(test, batch_size=64, shuffle=False)

# 4. Initialize the model: inputs, label ("mortality"), and mode ("binary")
#    all come from the task schema, so only hyperparameters are passed
model = Transformer(dataset=sample_dataset, embedding_dim=128)

trainer = Trainer(model=model, metrics=["pr_auc", "roc_auc", "f1"])  # device auto-detected
trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    monitor="pr_auc",            # AUPRC — robust for the rare-mortality class
    monitor_criterion="max",
)

# 5. Evaluate (uses the metrics passed to the Trainer)
results = trainer.evaluate(test_loader)
```

## Detailed Documentation

Read the reference file that matches the step you are on:

| File | Read when | Key topics |
|------|-----------|------------|
| `references/datasets.md` | Loading MIMIC/eICU/OMOP/signal datasets, splitting data | Patient/Event structures, loaders, `split_by_patient` / `split_by_visit` / `split_by_sample` |
| `references/medical_coding.md` | Translating or grouping ICD, NDC, RxNorm, ATC, CCS codes | `InnerMap` lookups and hierarchy, `CrossMap` translation |
| `references/tasks.md` | Choosing a predefined task or writing a custom one | 2.x task classes, `input_schema` / `output_schema`, custom `BaseTask` |
| `references/models.md` | Selecting and configuring a model | Baselines, RNN/CNN/Transformer, RETAIN, SafeDrug, GAMENet, StageNet, GAT/GCN |
| `references/preprocessing.md` | Understanding how raw events become tensors | Schema-string processors (`"sequence"`, `"timeseries"`, `"binary"`, ...) |
| `references/training_evaluation.md` | Training, metrics, calibration, uncertainty, interpretability | `Trainer`, metric strings, conformal prediction, Chefer/IG attributions |

## Installation

```bash
uv venv --python 3.13 .venv-pyhealth      # PyHealth 2.0.2 supports Python 3.12-3.13
source .venv-pyhealth/bin/activate
uv pip install "pyhealth>=2.0.2"
```

**Requirements (PyHealth 2.0.2):**
- Python **3.12 or 3.13** (`>=3.12,<3.14`) — if your default interpreter is 3.14, create the environment with `--python 3.13`.
- PyTorch, polars, pandas, scikit-learn, and transformers are installed as pinned dependencies — keep PyHealth in its own environment so these pins don't collide with other projects.
- A `2.1` alpha line exists on PyPI; stay on the 2.0.x releases unless you need an alpha-only feature.

## Common Use Cases

### Use Case 1: ICU Mortality Prediction

**Objective**: Predict patient mortality in intensive care unit

**Approach:**
1. Load MIMIC-IV dataset → Read `references/datasets.md`
2. Apply mortality prediction task → Read `references/tasks.md`
3. Select an interpretable model (RETAIN) or an attribution-friendly one (Transformer) → Read `references/models.md`
4. Train and evaluate → Read `references/training_evaluation.md`
5. Interpret predictions for clinical review → Read `references/training_evaluation.md`

### Use Case 2: Safe Medication Recommendation

**Objective**: Recommend medications while avoiding drug-drug interactions

**Approach:**
1. Load EHR dataset (MIMIC-III/IV, eICU, or OMOP) → Read `references/datasets.md`
2. Apply a `DrugRecommendation*` task → Read `references/tasks.md`
3. Use SafeDrug or GAMENet, which build their DDI graphs from the dataset → Read `references/models.md`
4. Preprocess medication codes → Read `references/medical_coding.md`
5. Evaluate with multi-label metrics (`jaccard_samples`, `f1_samples`, `pr_auc_samples`) → Read `references/training_evaluation.md`

### Use Case 3: Hospital Readmission Prediction

**Objective**: Identify patients at risk of readmission

**Approach:**
1. Load multi-site EHR data (eICU or OMOP) → Read `references/datasets.md`
2. Apply a `ReadmissionPrediction*` task → Read `references/tasks.md`
3. Handle class imbalance (report AUPRC, not only AUROC) → Read `references/training_evaluation.md`
4. Train a Transformer or RNN model → Read `references/models.md`
5. Calibrate predictions and assess fairness → Read `references/training_evaluation.md`

### Use Case 4: Sleep Staging

**Objective**: Classify sleep stages from EEG signals

**Approach:**
1. Load a sleep EEG dataset (SleepEDF, SHHS, ISRUC) → Read `references/datasets.md`
2. Apply sleep staging (`SleepStagingSleepEDF` or the legacy `sleep_staging_*_fn` functions) → Read `references/tasks.md`
3. Preprocess EEG signals (filtering, segmentation) → Read `references/preprocessing.md`
4. Train a CNN, SparcNet, or ContraWR model → Read `references/models.md`
5. Evaluate per-stage performance (`f1_macro`, `cohen_kappa`) → Read `references/training_evaluation.md`

### Use Case 5: Medical Code Translation

**Objective**: Standardize diagnoses across different coding systems

**Approach:**
1. Read `references/medical_coding.md` for comprehensive guidance
2. Use `CrossMap` to translate between ICD-9, ICD-10, and CCS
3. Group codes into clinically meaningful categories
4. Integrate with dataset processing

### Use Case 6: Clinical Text to ICD Coding

**Objective**: Automatically assign ICD codes from clinical notes

**Approach:**
1. Load MIMIC-III with clinical notes → Read `references/datasets.md`
2. Apply the `MIMIC3ICD9Coding` task → Read `references/tasks.md`
3. Preprocess clinical text → Read `references/preprocessing.md`
4. Use `TransformersModel(dataset=..., model_name="emilyalsentzer/Bio_ClinicalBERT")` → Read `references/models.md`
5. Evaluate with multi-label metrics → Read `references/training_evaluation.md`

## Best Practices

### Data Handling

1. **Always split by patient**: Prevent data leakage by ensuring no patient appears in multiple splits
   ```python
   from pyhealth.datasets import split_by_patient
   train, val, test = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=42)
   ```

2. **Check dataset statistics**: Understand your data before modeling
   ```python
   dataset.stats()  # prints patient and event counts (returns None)
   ```

3. **Use appropriate preprocessing**: Match processors to data types (see `references/preprocessing.md`)

### Model Development

1. **Start with baselines**: Establish baseline performance with simple models
   - `LogisticRegression` for binary/multi-class tasks
   - `MLP` for an initial deep learning baseline

2. **Choose task-appropriate models**:
   - Interpretability needed → RETAIN, AdaCare (by design); Transformer (post-hoc attributions)
   - Drug recommendation → SafeDrug, GAMENet
   - Long sequences → Transformer
   - Graph relationships → GAT / GCN

3. **Monitor validation metrics**: Use appropriate metrics for the task and handle class imbalance. PyHealth metric strings (pass to `Trainer(metrics=[...])` / `monitor=`):
   - Binary: `roc_auc`, `pr_auc` (prefer `pr_auc` for rare events), `f1`, `accuracy`
   - Multi-class: `f1_macro`, `f1_weighted`, `accuracy`, `cohen_kappa`
   - Multi-label / drug-rec: `jaccard_samples`, `f1_samples`, `pr_auc_samples`, `ddi` (reported as `ddi_score`)
   - Regression: `mae`, `mse`, `kl_divergence`

### Clinical Validation

1. **Calibrate predictions**: Ensure probabilities are reliable (see `references/training_evaluation.md`)
2. **Assess fairness**: Evaluate across demographic groups to detect bias
3. **Quantify uncertainty**: Provide confidence estimates for predictions (conformal prediction sets)
4. **Interpret predictions**: Attention/relevance maps, SHAP, or integrated gradients for clinician review
5. **Validate thoroughly**: Use held-out test sets from different time periods or sites
6. **Report transparently**: Follow TRIPOD+AI (BMJ 2024;385:e078378) when publishing a clinical prediction model

## Limitations and Considerations

### Data Requirements

- **Large datasets**: Deep learning models require sufficient data (thousands of patients)
- **Data quality**: Missing data and coding errors impact performance
- **Temporal consistency**: Ensure train/test split respects temporal ordering when needed
- **Access**: MIMIC and eICU require PhysioNet credentialing and a data use agreement; never copy restricted records into prompts, notebooks, or repositories that the agreement does not cover

### Clinical Validation

- **External validation**: Test on data from different hospitals/systems
- **Prospective evaluation**: Validate in real clinical settings before deployment
- **Clinical review**: Have clinicians review predictions and interpretations
- **Decision support, not diagnosis**: Present model outputs as research-grade risk estimates for qualified clinicians; deployment as a medical device falls under device regulation (e.g. FDA SaMD, EU MDR)
- **Ethical considerations**: Address privacy (HIPAA/GDPR), fairness, and safety

### Computational Resources

- **GPU recommended**: For training deep learning models efficiently
- **Memory requirements**: Large datasets may require 16GB+ RAM
- **Storage**: Healthcare datasets can be 10s-100s of GB, plus the task-sample cache

## Troubleshooting

### Common Issues

**`TypeError: ... unexpected keyword argument 'feature_keys'` (or `'root'`)**:
- You are using a 1.x-style call. Pass only `dataset=` and hyperparameters to models; use `MIMIC4EHRDataset(root=..., tables=...)` or `MIMIC4Dataset(ehr_root=..., ehr_tables=...)`

**ImportError or missing tables**:
- Ensure dataset files are downloaded and the root path points at the versioned folder
- Confirm the table names exist in the dataset's YAML config

**Out of memory**:
- Reduce batch size
- Reduce sequence length (`max_seq_len` on `Transformer`)
- Pass `dev=True` to the dataset loader (e.g. `MIMIC4EHRDataset(..., dev=True)`) to prototype on the first 1,000 patients
- Process data in chunks

**Poor performance**:
- Check class imbalance and use appropriate metrics (`pr_auc` vs `roc_auc`)
- Verify preprocessing (normalization, missing data handling)
- Increase model capacity or training epochs
- Check for data leakage in the train/test split

**Slow training**:
- Use a GPU (`Trainer(..., device="cuda")`)
- Increase batch size (if memory allows)
- Reduce sequence length
- Use a lighter model (CNN or RNN instead of Transformer)

### Getting Help

- **Documentation**: https://pyhealth.readthedocs.io/
- **GitHub Issues**: https://github.com/sunlabuiuc/PyHealth/issues
- **Examples/notebooks**: https://github.com/sunlabuiuc/PyHealth/tree/master/examples

## Example: Complete Workflow

```python
# Complete mortality prediction pipeline (PyHealth 2.0.x)
import torch
from pyhealth.datasets import MIMIC4EHRDataset, split_by_patient, get_dataloader
from pyhealth.tasks import MortalityPredictionMIMIC4
from pyhealth.models import Transformer
from pyhealth.trainer import Trainer
from pyhealth.interpret.methods import CheferRelevance

# 1. Load dataset (declare the tables the task needs)
dataset = MIMIC4EHRDataset(
    root="/data/mimic-iv/2.2",
    tables=["diagnoses_icd", "procedures_icd", "prescriptions"],
)
dataset.stats()

# 2. Define task (instantiate the task class)
sample_dataset = dataset.set_task(MortalityPredictionMIMIC4())
print(f"Generated {len(sample_dataset)} samples")

# 3. Split data (by patient to prevent leakage)
train_ds, val_ds, test_ds = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=42)

# 4. Create data loaders
train_loader = get_dataloader(train_ds, batch_size=64, shuffle=True)
val_loader = get_dataloader(val_ds, batch_size=64)
test_loader = get_dataloader(test_ds, batch_size=64)

# 5. Initialize the model (schema-driven; swap in RETAIN(dataset=sample_dataset,
#    embedding_dim=128) for a model that is interpretable by design)
model = Transformer(dataset=sample_dataset, embedding_dim=128, heads=2, num_layers=2)

# 6. Train, keeping the best checkpoint by validation AUPRC
trainer = Trainer(model=model, metrics=["accuracy", "pr_auc", "roc_auc", "f1"])
trainer.train(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    epochs=50,
    optimizer_class=torch.optim.Adam,
    optimizer_params={"lr": 1e-3},
    weight_decay=1e-5,
    monitor="pr_auc",          # AUPRC for the imbalanced (rare-mortality) outcome
    monitor_criterion="max",
    patience=5,                # early stopping
)

# 7. Evaluate on the test set (uses the metrics passed to the Trainer)
for metric, value in trainer.evaluate(test_loader).items():
    print(f"  {metric}: {value:.4f}")

# 8. Predictions with patient IDs: inference() returns (y_true, y_prob, loss),
#    extended with patient_ids when return_patient_ids=True
y_true, y_prob, loss, patient_ids = trainer.inference(test_loader, return_patient_ids=True)
positive_prob = y_prob if y_prob.ndim == 1 else y_prob[..., -1]
high_risk_idx = int(positive_prob.argmax())
print(f"Highest-risk patient: {patient_ids[high_risk_idx]} ({float(positive_prob[high_risk_idx]):.3f})")

# 9. Token-level relevance (Chefer; supported by Transformer and StageAttentionNet)
relevance = CheferRelevance(model)
batch = next(iter(get_dataloader(test_ds, batch_size=1, shuffle=False)))
for feature_key, rel in relevance.attribute(**batch).items():
    print(f"{feature_key}: top tokens -> {rel[0].topk(min(5, rel.shape[-1])).indices.tolist()}")

# 10. Save the trained weights
trainer.save_ckpt("./models/mortality_transformer.pt")
```

## Resources

For detailed information on each component, see the reference files in `references/`: `datasets.md`, `medical_coding.md`, `tasks.md`, `models.md`, `preprocessing.md`, and `training_evaluation.md` (see the table under **Detailed Documentation** for when to read each).

Part of the AlterLab Academic Skills suite.