JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jiarui Li
2026-06-13 15:39:31 +08:00
parent 46a3dfe628
commit 76787d2fb2
1 changed files with 124 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
evaluate_doa_auc.py
View File

@@ -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,12 +603,21 @@ 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))
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_ids,
disease_ids=disease_chunk,
horizons=horizons,
dist_mode=dist_mode,
score_mode=score_mode,
@@ -526,12 +626,15 @@ def main() -> None:
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}")