From e040707d0713492c14f1f5efa1eb7a18e86dc4cf Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Mon, 3 Aug 2026 13:56:06 +0800 Subject: [PATCH] fix: evaluate AUC on fixed test EIDs --- eval_data.py | 33 ++++++++ evaluate_auc.py | 59 +++++++++---- evaluate_auc_v2.py | 52 +++++++++--- tests/test_auc_eid_split.py | 162 ++++++++++++++++++++++++++++++++++++ 4 files changed, 280 insertions(+), 26 deletions(-) create mode 100644 tests/test_auc_eid_split.py diff --git a/eval_data.py b/eval_data.py index c671e97..4c598e2 100644 --- a/eval_data.py +++ b/eval_data.py @@ -135,6 +135,39 @@ def split_indices( ) +def select_indices_by_eid_file( + dataset: Any, + eid_file: str | Path, +) -> Tuple[np.ndarray, Path]: + """Return dataset indices whose patient EIDs occur in ``eid_file``.""" + from train_util import load_eid_file + + path = Path(eid_file) + if not path.is_absolute(): + direct = Path.cwd() / path + path = direct if direct.is_file() else Path(__file__).resolve().parent / path + if not path.is_file(): + raise FileNotFoundError(f"EID split file not found: {path}") + + samples = getattr(dataset, "samples", None) + if samples is None: + raise TypeError("EID-based evaluation requires dataset.samples") + selected_eids = load_eid_file(path) + indices = np.asarray( + [ + index + for index, sample in enumerate(samples) + if int(sample["eid"]) in selected_eids + ], + dtype=np.int64, + ) + if indices.size == 0: + raise ValueError( + f"No dataset patients matched the EID split file: {path}" + ) + return indices, path.resolve() + + def build_model_from_dataset( args: argparse.Namespace, cfg: Dict[str, Any], diff --git a/evaluate_auc.py b/evaluate_auc.py index 5e18b4d..c1510a9 100644 --- a/evaluate_auc.py +++ b/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, diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 1a1ece1..e1d7459 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -47,6 +47,7 @@ from eval_data import ( load_json_config, load_sequence_eval_dataset, resolve_eval_device, + select_indices_by_eid_file, split_indices, validate_training_mode_config, validate_dataset_metadata, @@ -96,7 +97,11 @@ def parse_float_list(value: Any) -> Optional[List[float]]: return [float(x.strip()) for x in text.split(",") if x.strip()] -def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dict[str, Any]) -> np.ndarray: +def make_eval_indices( + dataset: HealthDataset, + args: argparse.Namespace, + cfg: Dict[str, Any], +) -> 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)) @@ -105,18 +110,32 @@ def make_eval_indices(dataset: HealthDataset, args: argparse.Namespace, cfg: Dic if eval_split in {"valid", "validation"}: eval_split = "val" - 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), dtype=np.int64), - } - if eval_split not in split_map: + if eval_split not in {"train", "val", "test", "all"}: raise ValueError(f"Unsupported eval_split={eval_split!r}") - indices = split_map[eval_split] + + 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), dtype=np.int64), + } + indices = split_map[eval_split] subset_size = cfg_get(args, cfg, "dataset_subset_size", None) if subset_size is not None and int(subset_size) > 0: @@ -1054,6 +1073,15 @@ def main() -> None: parser.add_argument("--output_path", type=str, default=None) parser.add_argument("--eval_split", type=str, default="test", choices=["train", "val", "valid", "validation", "test", "all"]) + 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) parser.add_argument("--batch_size", type=int, default=None) diff --git a/tests/test_auc_eid_split.py b/tests/test_auc_eid_split.py new file mode 100644 index 0000000..a230548 --- /dev/null +++ b/tests/test_auc_eid_split.py @@ -0,0 +1,162 @@ +import argparse +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np + +from eval_data import split_indices +from evaluate_auc import make_eval_subset +from evaluate_auc_v2 import make_eval_indices + + +class _DummyDataset: + def __init__(self, eids: list[int]) -> None: + self.samples = [{"eid": eid} for eid in eids] + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, index: int) -> dict[str, int]: + return self.samples[index] + + +class AUCEidSplitTests(unittest.TestCase): + def test_both_auc_evaluators_default_to_ukb_test_eid_file(self) -> None: + dataset = _DummyDataset([101, 102, 103, 104, 105]) + args = argparse.Namespace( + eval_split="test", + dataset_subset_size=None, + test_eid_file=None, + ) + + with tempfile.TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + (root / "ukb_test_eid.csv").write_text( + "eid\n104\n102\n", + encoding="utf-8", + ) + with patch.object(Path, "cwd", return_value=root): + subset, legacy_indices = make_eval_subset(dataset, args, {}) + landmark_indices = make_eval_indices(dataset, args, {}) + + expected = np.asarray([1, 3], dtype=np.int64) + np.testing.assert_array_equal(legacy_indices, expected) + np.testing.assert_array_equal(landmark_indices, expected) + self.assertEqual(subset.indices, expected.tolist()) + + def test_subset_size_is_applied_after_eid_selection(self) -> None: + dataset = _DummyDataset([201, 202, 203, 204]) + args = argparse.Namespace( + eval_split="test", + dataset_subset_size=1, + test_eid_file=None, + ) + + with tempfile.TemporaryDirectory() as tmp_dir: + eid_path = Path(tmp_dir) / "test.csv" + eid_path.write_text("eid\n202\n204\n", encoding="utf-8") + cfg = {"test_eid_file": str(eid_path)} + _, legacy_indices = make_eval_subset(dataset, args, cfg) + landmark_indices = make_eval_indices(dataset, args, cfg) + + expected = np.asarray([1], dtype=np.int64) + np.testing.assert_array_equal(legacy_indices, expected) + np.testing.assert_array_equal(landmark_indices, expected) + + def test_cli_test_eid_file_overrides_config(self) -> None: + dataset = _DummyDataset([301, 302, 303]) + + with tempfile.TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + config_path = root / "config.csv" + cli_path = root / "cli.csv" + config_path.write_text("eid\n301\n", encoding="utf-8") + cli_path.write_text("eid\n303\n", encoding="utf-8") + args = argparse.Namespace( + eval_split="test", + dataset_subset_size=None, + test_eid_file=str(cli_path), + ) + cfg = {"test_eid_file": str(config_path)} + + _, legacy_indices = make_eval_subset(dataset, args, cfg) + landmark_indices = make_eval_indices(dataset, args, cfg) + + expected = np.asarray([2], dtype=np.int64) + np.testing.assert_array_equal(legacy_indices, expected) + np.testing.assert_array_equal(landmark_indices, expected) + + def test_empty_test_eid_file_explicitly_uses_ratio_split(self) -> None: + dataset = _DummyDataset(list(range(20))) + args = argparse.Namespace( + eval_split="test", + dataset_subset_size=None, + test_eid_file="", + ) + cfg = { + "train_ratio": 0.7, + "val_ratio": 0.15, + "test_ratio": 0.15, + "seed": 7, + } + + expected = split_indices(20, 0.7, 0.15, 0.15, 7)[2] + _, legacy_indices = make_eval_subset(dataset, args, cfg) + landmark_indices = make_eval_indices(dataset, args, cfg) + + np.testing.assert_array_equal(legacy_indices, expected) + np.testing.assert_array_equal(landmark_indices, expected) + + def test_non_test_split_does_not_read_test_eid_file(self) -> None: + dataset = _DummyDataset(list(range(20))) + args = argparse.Namespace( + eval_split="val", + dataset_subset_size=None, + test_eid_file="missing.csv", + ) + cfg = { + "train_ratio": 0.7, + "val_ratio": 0.15, + "test_ratio": 0.15, + "seed": 11, + } + + expected = split_indices(20, 0.7, 0.15, 0.15, 11)[1] + _, legacy_indices = make_eval_subset(dataset, args, cfg) + landmark_indices = make_eval_indices(dataset, args, cfg) + + np.testing.assert_array_equal(legacy_indices, expected) + np.testing.assert_array_equal(landmark_indices, expected) + + def test_missing_or_nonmatching_eid_file_fails_closed(self) -> None: + dataset = _DummyDataset([401, 402]) + args = argparse.Namespace( + eval_split="test", + dataset_subset_size=None, + test_eid_file=None, + ) + + with tempfile.TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + missing = root / "missing.csv" + with self.assertRaisesRegex(FileNotFoundError, "EID split file"): + make_eval_indices( + dataset, + args, + {"test_eid_file": str(missing)}, + ) + + nonmatching = root / "nonmatching.csv" + nonmatching.write_text("eid\n999\n", encoding="utf-8") + with self.assertRaisesRegex(ValueError, "No dataset patients"): + make_eval_indices( + dataset, + args, + {"test_eid_file": str(nonmatching)}, + ) + + +if __name__ == "__main__": + unittest.main()