diff --git a/evaluate_doa_auc.py b/evaluate_doa_auc.py index f30b7e8..aa74076 100644 --- a/evaluate_doa_auc.py +++ b/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 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( - dataset=dataset, - hidden_all=hidden_all, - row_arrays=row_arrays, - model=model, - disease_ids=disease_ids, - 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)), - ) + chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 256)) + result_parts = [] + for disease_chunk in tqdm( + list(iter_chunks(disease_ids, chunk_size)), + desc="Disease chunks", + leave=True, + dynamic_ncols=True, + ): + result_parts.append( + evaluate_doa_auc( + dataset=dataset, + hidden_all=hidden_all, + 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}")