sunlabuiuc · jhnwu3 · Feb 17, 2026 · Feb 16, 2026 · Feb 16, 2026 · Feb 16, 2026
diff --git a/docs/api/graph.rst b/docs/api/graph.rst
diff --git a/docs/api/graph/pyhealth.graph.KnowledgeGraph.rst b/docs/api/graph/pyhealth.graph.KnowledgeGraph.rst
@@ -0,0 +1,15 @@
+pyhealth.graph.KnowledgeGraph
+==============================
+
+Overview
+--------
+Knowledge graph data structure for healthcare code systems.
+Stores (head, relation, tail) triples and provides k-hop subgraph
+extraction for patient-level graph construction.
+
+API Reference
+-------------
+.. automodule:: pyhealth.graph.knowledge_graph
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/api/processors.rst b/docs/api/processors.rst
@@ -47,6 +47,7 @@ Available Processors
 - ``StageNetTensorProcessor``: Tensor processing for StageNet
 - ``MultiHotProcessor``: For multi-hot encoding
 - ``IgnoreProcessor``: A special feature processor that marks a feature to be ignored.
+- ``GraphProcessor``: For knowledge graph subgraph extraction (e.g., GraphCare, G-BERT)
 
 Usage Examples
 --------------
@@ -276,6 +277,7 @@ Common string keys for automatic processor selection:
 - ``"time_image"``: For time-stamped image sequences
 - ``"tensor"``: For pre-processed tensors
 - ``"raw"``: For raw/unprocessed data
+- ``"graph"``: For knowledge graph subgraphs
 
 Writing Custom FeatureProcessors
 ---------------------------------
@@ -470,4 +472,5 @@ API Reference
     processors/pyhealth.processors.IgnoreProcessor
     processors/pyhealth.processors.MultiHotProcessor
     processors/pyhealth.processors.StageNetProcessor
-    processors/pyhealth.processors.StageNetTensorProcessor
+    processors/pyhealth.processors.StageNetTensorProcessor
+    processors/pyhealth.processors.GraphProcessor
diff --git a/docs/api/processors/pyhealth.processors.GraphProcessor.rst b/docs/api/processors/pyhealth.processors.GraphProcessor.rst
@@ -0,0 +1,15 @@
+pyhealth.processors.GraphProcessor
+====================================
+
+Overview
+--------
+Processor that converts medical codes into patient-level PyG subgraphs
+using a provided KnowledgeGraph. Registered as ``"graph"`` in the
+processor registry.
+
+API Reference
+-------------
+.. automodule:: pyhealth.processors.graph_processor
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/index.rst b/docs/index.rst
@@ -223,6 +223,7 @@ Quick Navigation
 
    api/data
    api/datasets
+   api/graph
    api/tasks
    api/models
    api/processors

diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py
@@ -15,6 +15,12 @@
 
 MODULE_CACHE_PATH = os.path.join(BASE_CACHE_PATH, "datasets")
 create_directory(MODULE_CACHE_PATH)
+#PyG import for graph-based models
+try:
+    from torch_geometric.data import Data as PyGData, Batch as PyGBatch
+    HAS_PYG = True
+except ImportError:
+    HAS_PYG = False
 
 
 # basic tables which are a part of the defined datasets
@@ -294,6 +300,9 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
 
             # Return as tuple (time, values)
             collated[key] = (collated_times, collated_values)
+            # PyG Data objects (graph processor output)
+        elif HAS_PYG and isinstance(values[0], PyGData):
+            collated[key] = PyGBatch.from_data_list(values)
 
         elif isinstance(values[0], torch.Tensor):
             # Check if shapes are the same

diff --git a/pyhealth/graph/__init__.py b/pyhealth/graph/__init__.py
@@ -0,0 +1,3 @@
+from pyhealth.graph.knowledge_graph import KnowledgeGraph
+
+__all__ = ["KnowledgeGraph"]
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from pyhealth.graph.knowledge_graph import KnowledgeGraph

		__all__ = ["KnowledgeGraph"]