Refactor code structure for improved readability and maintainability

train_util.py

@@ -4,6 +4,7 @@ import json
 import logging
 import sys
 import time
+import csv
 from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, Iterable, Tuple
@@ -86,6 +87,26 @@ def load_extra_info_types_file(path: str) -> list[int]:
         raise ValueError(f"Invalid extra info type id in {path}") from exc
 
 
+def load_eid_file(path: str | Path) -> set[int]:
+    file_path = Path(path)
+    if not file_path.is_file():
+        raise FileNotFoundError(f"eid split file not found: {file_path}")
+    with file_path.open(newline="", encoding="utf-8-sig") as f:
+        reader = csv.DictReader(f)
+        if reader.fieldnames is None or "eid" not in reader.fieldnames:
+            raise ValueError(
+                f"eid split file must contain an 'eid' column: {file_path}"
+            )
+        out: set[int] = set()
+        for row in reader:
+            raw = (row.get("eid") or "").strip()
+            if raw:
+                out.add(int(raw))
+    if not out:
+        raise ValueError(f"eid split file is empty: {file_path}")
+    return out
+
+
 def configure_torch_for_training(device: torch.device) -> None:
     if device.type != "cuda":
         return
@@ -132,6 +153,44 @@ def split_dataset(
     )
 
 
+def split_dataset_by_eid_files(
+    dataset: HealthDataset,
+    train_eid_file: str | Path,
+    val_eid_file: str | Path,
+    test_eid_file: str | Path,
+) -> Tuple[Subset, Subset, Subset]:
+    split_sets = {
+        "train": load_eid_file(train_eid_file),
+        "val": load_eid_file(val_eid_file),
+        "test": load_eid_file(test_eid_file),
+    }
+    overlaps = (
+        split_sets["train"] & split_sets["val"],
+        split_sets["train"] & split_sets["test"],
+        split_sets["val"] & split_sets["test"],
+    )
+    if any(overlaps):
+        raise ValueError("eid split files must be disjoint")
+
+    split_indices: Dict[str, list[int]] = {"train": [], "val": [], "test": []}
+    for idx, sample in enumerate(dataset.samples):
+        eid = int(sample["eid"])
+        for split_name, eid_set in split_sets.items():
+            if eid in eid_set:
+                split_indices[split_name].append(idx)
+                break
+
+    missing = [name for name, indices in split_indices.items() if not indices]
+    if missing:
+        raise ValueError(f"Empty dataset split(s) after eid filtering: {missing}")
+
+    return (
+        Subset(dataset, np.asarray(split_indices["train"], dtype=np.int64)),
+        Subset(dataset, np.asarray(split_indices["val"], dtype=np.int64)),
+        Subset(dataset, np.asarray(split_indices["test"], dtype=np.int64)),
+    )
+
+
 def split_all_future_datasets(
     train_dataset: AllFutureHealthDataset,
     val_dataset: AllFutureHealthDataset,
@@ -172,6 +231,57 @@ def split_all_future_datasets(
     )
 
 
+def split_all_future_datasets_by_eid_files(
+    train_dataset: AllFutureHealthDataset,
+    val_dataset: AllFutureHealthDataset,
+    test_dataset: AllFutureHealthDataset,
+    train_eid_file: str | Path,
+    val_eid_file: str | Path,
+    test_eid_file: str | Path,
+) -> Tuple[Subset, Subset, Subset]:
+    split_sets = {
+        "train": load_eid_file(train_eid_file),
+        "val": load_eid_file(val_eid_file),
+        "test": load_eid_file(test_eid_file),
+    }
+    overlaps = (
+        split_sets["train"] & split_sets["val"],
+        split_sets["train"] & split_sets["test"],
+        split_sets["val"] & split_sets["test"],
+    )
+    if any(overlaps):
+        raise ValueError("eid split files must be disjoint")
+
+    train_patient_idx = [
+        idx
+        for idx, patient in enumerate(train_dataset.patients)
+        if int(patient["eid"]) in split_sets["train"]
+    ]
+    val_query_idx = [
+        idx
+        for idx, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
+        if int(val_dataset.patients[int(pidx)]["eid"]) in split_sets["val"]
+    ]
+    test_query_idx = [
+        idx
+        for idx, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
+        if int(test_dataset.patients[int(pidx)]["eid"]) in split_sets["test"]
+    ]
+
+    if not train_patient_idx:
+        raise ValueError("All-future training eid split has no patients.")
+    if not val_query_idx:
+        raise ValueError("All-future validation eid split has no valid query samples.")
+    if not test_query_idx:
+        raise ValueError("All-future test eid split has no valid query samples.")
+
+    return (
+        Subset(train_dataset, np.asarray(train_patient_idx, dtype=np.int64)),
+        Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)),
+        Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)),
+    )
+
+
 def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
     return AdamW(
         model.parameters(),