Refactor code structure for improved readability and maintainability

train_all_future.py

@@ -40,6 +40,7 @@ from train_util import (
     set_seed,
     setup_logging,
     split_all_future_datasets,
+    split_all_future_datasets_by_eid_files,
 )
 
 
@@ -55,6 +56,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--train_ratio", type=float, default=0.7)
     parser.add_argument("--val_ratio", type=float, default=0.15)
     parser.add_argument("--test_ratio", type=float, default=0.15)
+    parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
+    parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
+    parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
     parser.add_argument("--min_history_events", type=int, default=1)
     parser.add_argument("--min_future_events", type=int, default=1)
     parser.add_argument("--validation_query_seed", type=int, default=None)
@@ -89,7 +93,11 @@ def parse_args() -> argparse.Namespace:
         raise ValueError("min_history_events must be >= 1")
     if args.min_future_events < 1:
         raise ValueError("min_future_events must be >= 1")
-    if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
+    use_eid_split = all(
+        getattr(args, name)
+        for name in ("train_eid_file", "val_eid_file", "test_eid_file")
+    )
+    if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
         raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
     if args.validation_query_seed is None:
         args.validation_query_seed = int(args.seed)
@@ -335,15 +343,33 @@ def main() -> None:
         validation_query_seed=args.validation_query_seed,
         extra_info_types=args.extra_info_types,
     )
-    train_subset, val_subset, test_subset = split_all_future_datasets(
-        train_dataset=train_dataset,
-        val_dataset=val_dataset,
-        test_dataset=test_dataset,
-        train_ratio=args.train_ratio,
-        val_ratio=args.val_ratio,
-        test_ratio=args.test_ratio,
-        seed=args.seed,
-    )
+    if args.train_eid_file and args.val_eid_file and args.test_eid_file:
+        train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            test_dataset=test_dataset,
+            train_eid_file=args.train_eid_file,
+            val_eid_file=args.val_eid_file,
+            test_eid_file=args.test_eid_file,
+        )
+        logger.info(
+            "Using eid split files: "
+            f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
+        )
+    else:
+        train_subset, val_subset, test_subset = split_all_future_datasets(
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            test_dataset=test_dataset,
+            train_ratio=args.train_ratio,
+            val_ratio=args.val_ratio,
+            test_ratio=args.test_ratio,
+            seed=args.seed,
+        )
+        logger.info(
+            f"Using random ratio split: train={args.train_ratio}, "
+            f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
+        )
     logger.info(
         f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
     )

train_next_step.py

@@ -38,6 +38,7 @@ from train_util import (
     set_seed,
     setup_logging,
     split_dataset,
+    split_dataset_by_eid_files,
 )
 
 
@@ -55,6 +56,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--train_ratio", type=float, default=0.7)
     parser.add_argument("--val_ratio", type=float, default=0.15)
     parser.add_argument("--test_ratio", type=float, default=0.15)
+    parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
+    parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
+    parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
 
     parser.add_argument("--n_embd", type=int, default=120)
     parser.add_argument("--n_head", type=int, default=10)
@@ -92,7 +96,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--device", type=str, default="cuda")
 
     args = parser.parse_args()
-    if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
+    use_eid_split = all(
+        getattr(args, name)
+        for name in ("train_eid_file", "val_eid_file", "test_eid_file")
+    )
+    if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
         raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
     if args.target_mode == "uts":
         args.readout_name = args.readout_name or "same_time_group_end"
@@ -477,13 +485,29 @@ def main() -> None:
         include_no_event_in_uts_target=args.include_no_event_in_uts_target,
         extra_info_types=args.extra_info_types,
     )
-    train_subset, val_subset, test_subset = split_dataset(
-        dataset=dataset,
-        train_ratio=args.train_ratio,
-        val_ratio=args.val_ratio,
-        test_ratio=args.test_ratio,
-        seed=args.seed,
-    )
+    if args.train_eid_file and args.val_eid_file and args.test_eid_file:
+        train_subset, val_subset, test_subset = split_dataset_by_eid_files(
+            dataset=dataset,
+            train_eid_file=args.train_eid_file,
+            val_eid_file=args.val_eid_file,
+            test_eid_file=args.test_eid_file,
+        )
+        logger.info(
+            "Using eid split files: "
+            f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
+        )
+    else:
+        train_subset, val_subset, test_subset = split_dataset(
+            dataset=dataset,
+            train_ratio=args.train_ratio,
+            val_ratio=args.val_ratio,
+            test_ratio=args.test_ratio,
+            seed=args.seed,
+        )
+        logger.info(
+            f"Using random ratio split: train={args.train_ratio}, "
+            f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
+        )
     logger.info(
         f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
     )

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(),