Add evaluation split handling and dataset subset size to DOA AUC evaluation

evaluate_doa_auc.py

@@ -23,7 +23,7 @@ import numpy as np
 import pandas as pd
 import torch
 from torch.nn.utils.rnn import pad_sequence
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, Subset
 from tqdm.auto import tqdm
 
 from dataset import _ExpoBaseDataset
@@ -56,6 +56,71 @@ def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: A
     return cfg.get(name, default)
 
 
+def split_indices(
+    n: int,
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float,
+    seed: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    total = float(train_ratio) + float(val_ratio) + float(test_ratio)
+    if not np.isclose(total, 1.0, atol=1e-6):
+        raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
+    indices = np.random.RandomState(int(seed)).permutation(int(n))
+    n_train = int(n * train_ratio)
+    n_val = int(n * val_ratio)
+    return indices[:n_train], indices[n_train:n_train + n_val], indices[n_train + n_val:]
+
+
+def make_eval_indices(
+    dataset: Dataset,
+    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))
+    seed = int(cfg_get(args, cfg, "seed", 42))
+    eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
+    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:
+        raise ValueError(f"Unsupported eval_split={eval_split!r}")
+
+    indices = np.asarray(split_map[eval_split], dtype=np.int64)
+    subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
+    if subset_size is not None and int(subset_size) > 0:
+        indices = indices[: int(subset_size)]
+    return indices
+
+
+def subset_first_occurrence_map(
+    first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
+    selected_patient_ids: np.ndarray,
+) -> Dict[int, Tuple[np.ndarray, np.ndarray]]:
+    selected = set(int(x) for x in np.asarray(selected_patient_ids, dtype=np.int64).tolist())
+    out: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
+    for token, pairs in first_occurrence_by_token.items():
+        p, t = pairs
+        keep = np.array([int(x) in selected for x in p], dtype=bool)
+        if np.any(keep):
+            out[int(token)] = (
+                np.asarray(p, dtype=np.int32)[keep],
+                np.asarray(t, dtype=np.float32)[keep],
+            )
+    return out
+
+
 class DOAStatusDataset(_ExpoBaseDataset):
     def __init__(
         self,
@@ -406,13 +471,26 @@ def evaluate_doa_auc(
     return pd.DataFrame(rows)
 
 
+def iter_chunks(values: Sequence[int], chunk_size: int) -> Iterable[List[int]]:
+    values = [int(x) for x in values]
+    if chunk_size <= 0:
+        yield values
+        return
+    for start in range(0, len(values), chunk_size):
+        yield values[start:start + chunk_size]
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description="Evaluate DOA fixed-horizon disease AUC")
     parser.add_argument("--run_path", type=str, required=True)
     parser.add_argument("--output_path", type=str, default=None)
+    parser.add_argument("--eval_split", type=str, default=None,
+                        choices=["train", "val", "valid", "validation", "test", "all"])
+    parser.add_argument("--dataset_subset_size", type=int, default=None)
     parser.add_argument("--batch_size", type=int, default=None)
     parser.add_argument("--num_workers", type=int, default=None)
     parser.add_argument("--logit_batch_size", type=int, default=None)
+    parser.add_argument("--disease_chunk_size", type=int, default=None)
     parser.add_argument("--horizons", type=str, default=None)
     parser.add_argument("--score_mode", type=str, choices=["risk", "eta"], default=None)
     parser.add_argument("--filter_min_total", type=int, default=None)
@@ -446,6 +524,15 @@ def main() -> None:
         extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
     )
     validate_dataset_metadata(dataset, cfg)
+    eval_indices = make_eval_indices(dataset, args, cfg)
+    eval_patient_ids = np.asarray(
+        [dataset.records[int(i)]["patient_id"] for i in eval_indices],
+        dtype=np.int32,
+    )
+    eval_first_occurrence = subset_first_occurrence_map(
+        dataset.first_occurrence_by_token,
+        eval_patient_ids,
+    )
 
     disease_requested = parse_int_list(cfg_get(args, cfg, "diseases_of_interest", None))
     disease_ids = select_disease_tokens(
@@ -453,7 +540,7 @@ def main() -> None:
         labels_meta=labels_meta,
         requested_tokens=disease_requested,
         filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
-        first_occurrence_by_token=dataset.first_occurrence_by_token,
+        first_occurrence_by_token=eval_first_occurrence,
     )
     if not disease_ids:
         raise RuntimeError("No disease tokens selected after filtering.")
@@ -481,8 +568,9 @@ def main() -> None:
     ):
         raise RuntimeError("Next-token DOA evaluation requires <NO_EVENT> in the model vocabulary.")
 
+    eval_dataset = Subset(dataset, eval_indices)
     loader = DataLoader(
-        dataset,
+        eval_dataset,
         batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
         shuffle=False,
         collate_fn=collate_doa_fn,
@@ -496,7 +584,10 @@ def main() -> None:
     readout_name = str(cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token"))
     readout_reduce = str(cfg.get("readout_reduce", "mean"))
 
-    print(f"DOA records: {len(dataset)}")
+    eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
+    if eval_split in {"valid", "validation"}:
+        eval_split = "val"
+    print(f"DOA records: total={len(dataset)}, eval_{eval_split}={len(eval_dataset)}")
     print(f"Model target mode: {model_target_mode}")
     print(f"Dist mode: {dist_mode}")
     print(f"Score mode: {score_mode}")
@@ -512,26 +603,38 @@ def main() -> None:
         readout_reduce=readout_reduce,
         use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
     )
-    result = evaluate_doa_auc(
+    chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 256))
-        dataset=dataset,
+    result_parts = []
-        hidden_all=hidden_all,
+    for disease_chunk in tqdm(
-        row_arrays=row_arrays,
+        list(iter_chunks(disease_ids, chunk_size)),
-        model=model,
+        desc="Disease chunks",
-        disease_ids=disease_ids,
+        leave=True,
-        horizons=horizons,
+        dynamic_ncols=True,
-        dist_mode=dist_mode,
+    ):
-        score_mode=score_mode,
+        result_parts.append(
-        min_cases=min_cases,
+            evaluate_doa_auc(
-        device=device,
+                dataset=dataset,
-        logit_batch_size=int(cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
+                hidden_all=hidden_all,
-        use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
+                row_arrays=row_arrays,
-    )
+                model=model,
+                disease_ids=disease_chunk,
+                horizons=horizons,
+                dist_mode=dist_mode,
+                score_mode=score_mode,
+                min_cases=min_cases,
+                device=device,
+                logit_batch_size=int(cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
+                use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
+            )
+        )
+    result = pd.concat(result_parts, ignore_index=True) if result_parts else pd.DataFrame()
     if result.empty:
         raise RuntimeError("No DOA AUC rows produced. Check disease selection and min_cases.")
 
     meta = build_metadata_for_merge(dataset, labels_meta)
     result = result.merge(meta, on="token", how="left")
-    out_file = output_path / "doa_auc.csv"
+    result.insert(0, "eval_split", eval_split)
+    out_file = output_path / f"doa_auc_{eval_split}.csv"
     result.to_csv(out_file, index=False)
 
     summary = result.groupby(["token", "label_code", "horizon"], dropna=False, as_index=False).agg(
@@ -539,7 +642,8 @@ def main() -> None:
         n_case=("n_case", "sum"),
         n_control=("n_control", "sum"),
     )
-    summary.to_csv(output_path / "doa_auc_summary.csv", index=False)
+    summary.insert(0, "eval_split", eval_split)
+    summary.to_csv(output_path / f"doa_auc_{eval_split}_summary.csv", index=False)
     print(f"Wrote {out_file}")