fix: evaluate AUC on fixed test EIDs

This commit is contained in:
2026-08-03 13:56:06 +08:00
parent 271a678e73
commit e040707d07
4 changed files with 280 additions and 26 deletions

View File

@@ -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],

View File

@@ -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,

View File

@@ -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)

162
tests/test_auc_eid_split.py Normal file
View File

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