fix: evaluate AUC on fixed test EIDs

evaluate_auc.py

@@ -56,6 +56,7 @@ from eval_data import (
     load_json_config,
     load_sequence_eval_dataset,
     resolve_eval_device,
+    select_indices_by_eid_file,
     sequence_eval_collate_fn,
     split_indices,
     validate_training_mode_config,
@@ -279,7 +280,11 @@ def load_model_state(
     model.load_state_dict(state, strict=True)
 
 
-def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]:
+def make_eval_subset(
+    dataset: HealthDataset,
+    args: argparse.Namespace | Dict[str, Any] | None,
+    cfg: Dict[str, Any],
+) -> Tuple[Subset, np.ndarray]:
     train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
     val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
     test_ratio = float(cfg_get(args, cfg, "test_ratio", 0.15))
@@ -287,21 +292,38 @@ def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str
     eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
     dataset_subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
 
-    train_idx, val_idx, test_idx = split_indices(
-        len(dataset), train_ratio, val_ratio, test_ratio, seed)
-    split_map = {
-        "train": train_idx,
-        "val": val_idx,
-        "valid": val_idx,
-        "validation": val_idx,
-        "test": test_idx,
-        "all": np.arange(len(dataset)),
-    }
-    if eval_split not in split_map:
+    if eval_split in {"valid", "validation"}:
+        eval_split = "val"
+    if eval_split not in {"train", "val", "test", "all"}:
         raise ValueError(
-            f"eval_split must be one of {sorted(split_map)}, got {eval_split!r}")
+            "eval_split must be one of train/val/test/all, got "
+            f"{eval_split!r}"
+        )
+
+    test_eid_file = cfg_get(
+        args,
+        cfg,
+        "test_eid_file",
+        "ukb_test_eid.csv",
+    )
+    if eval_split == "test" and test_eid_file not in {None, ""}:
+        indices, eid_path = select_indices_by_eid_file(
+            dataset,
+            str(test_eid_file),
+        )
+        print(f"Test split source: EID file {eid_path}")
+    else:
+        train_idx, val_idx, test_idx = split_indices(
+            len(dataset), train_ratio, val_ratio, test_ratio, seed
+        )
+        split_map = {
+            "train": train_idx,
+            "val": val_idx,
+            "test": test_idx,
+            "all": np.arange(len(dataset)),
+        }
+        indices = split_map[eval_split]
 
-    indices = split_map[eval_split]
     if dataset_subset_size is not None and int(dataset_subset_size) > 0:
         indices = indices[: int(dataset_subset_size)]
     return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)
@@ -1129,6 +1151,15 @@ def main() -> None:
                         choices=["train", "val", "valid",
                                  "validation", "test", "all"],
                         help="Evaluation split. Defaults to 'test' unless cfg contains eval_split.")
+    parser.add_argument(
+        "--test_eid_file",
+        type=str,
+        default=None,
+        help=(
+            "Patient EID file for the test split. Defaults to train_config.json "
+            "or ukb_test_eid.csv. Set to an empty value to use ratio splitting."
+        ),
+    )
     parser.add_argument("--dataset_subset_size", type=int, default=None,
                         help="Optional number of patients from the selected split.")
     parser.add_argument("--batch_size", type=int, default=None,