[Contribution] GraphCare: KG-enhanced EHR predictions (ICLR 2024)

[joshuasteier]

Contributor Information

• Name: Joshua Steier
• Contribution Type: Model + Tests + Documentation

Description

Adds GraphCare (Jiang et al., ICLR 2024) as a native PyHealth model. GraphCare uses personalized knowledge graphs with a GNN backbone for healthcare prediction tasks.

Paper: [GraphCare: Enhancing Healthcare Predictions with Personalized Knowledge Graphs](https://openreview.net/forum?id=tVTN7Zs0ml)
Original repo: [pat-jj/GraphCare](https://github.com/pat-jj/GraphCare)

This PR builds directly on the pyhealth.graph infrastructure from #853, proving it works end-to-end with a real model.

What's included:

pyhealth.models.GraphCare — Full model inheriting BaseModel:

• Three GNN backbones: BAT (Bi-Attention GNN, from paper), GAT, GIN
• Three patient representation modes: joint, graph, node
• Node-level alpha attention with softmax pooling
• Optional edge relation weights
• Residual connections and layer normalization
• Supports pre-computed node features (TransE, LLM embeddings)
• Standard forward() returning {loss, y_prob, y_true, logit}

pyhealth.models._graphcare.BiAttentionGNNConv — BAT GNN layer:

• GIN-style message passing with node attention and edge relation weights
• Optional edge attention via learnable W_R projection

Usage

from pyhealth.graph import KnowledgeGraph
from pyhealth.models import GraphCare

kg = KnowledgeGraph(triples="data/umls_triples.csv")

input_schema = {
    "conditions": ("graph", {"knowledge_graph": kg, "num_hops": 2}),
    "procedures": ("graph", {"knowledge_graph": kg, "num_hops": 2}),
}

model = GraphCare(
    dataset=sample_dataset,
    knowledge_graph=kg,
    hidden_dim=128,
    num_layers=3,
    gnn_type="bat",
    patient_mode="joint",
)

Design decisions:

• Consumes pyhealth.graph.KnowledgeGraph and GraphProcessor from [Contribution] pyhealth.graph: Native knowledge graph support + GraphProcessor #853 — no custom KG builder or DataLoader needed
• torch-geometric remains optional — all imports guarded with try/except
• Deleted old kg_builder.py (superseded by pyhealth.graph)
• Follows same patterns as JambaEHR and other PyHealth models

Files to Review

New files:

• pyhealth/models/graphcare.py — main model
• tests/core/test_graphcare.py — 22 unit tests (BAT, GAT, GIN × 3 modes)
• docs/api/models/pyhealth.models.GraphCare.rst — API docs

Modified files:

• pyhealth/models/_graphcare/__init__.py — removed kg_builder import
• pyhealth/models/_graphcare/bat_gnn.py — optional PyG imports, string-quoted annotations
• pyhealth/models/__init__.py — added GraphCare export
• docs/api/models.rst — added to toctree

Deleted files:

• pyhealth/models/_graphcare/kg_builder.py — superseded by pyhealth.graph.KnowledgeGraph

Testing

python -m unittest tests.core.test_graphcare -v
# 22 tests (skip gracefully if torch-geometric not installed)

Dependencies

• torch-geometric (optional, same as [Contribution] pyhealth.graph: Native knowledge graph support + GraphProcessor #853)
• pyhealth.graph module ([Contribution] pyhealth.graph: Native knowledge graph support + GraphProcessor #853, already on master)

[jhnwu3]

Is there a way to make an example script using mortality prediction as an example?

Not sure if that's easy, let me know.

[joshuasteier]

@jhnwu3:
Done: added examples/graphcare_tutorial.ipynb with a
synthetic mortality prediction example. Covers all 3
backbones (BAT/GAT/GIN), all 3 patient modes
(joint/graph/node), and includes a Section 10 showing
how to use it with real MIMIC data via the existing
MortalityPredictionMIMIC3 task. Let me know what you think!

[jhnwu3]

Lgtm, actually I was thinking of saying we could probably redefine the MIMIC3 or MIMIC4 task to have a swappable input_schema, but this example actually highlights a nice flexibility of PyHealth I hadn't even considered before. That's really cool!