777 lines
30 KiB
Python
777 lines
30 KiB
Python
"""Evaluate disease AUC at date of assessment (DOA).
|
|
|
|
Cases are patients whose first occurrence of a disease is after DOA and within
|
|
the requested horizon. Controls are patients who never have that disease in the
|
|
full observed record. Patients prevalent at/before DOA or incident after the
|
|
horizon are not used for that disease-horizon AUC.
|
|
|
|
The script adapts automatically to checkpoint target mode:
|
|
- next_token: use the DOA token position, inserting <NO_EVENT> at DOA when no
|
|
real disease token exists at DOA;
|
|
- all_future: query the model directly with t_query=DOA, allowing empty
|
|
disease history because other-info tokens still describe the DOA state.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import os
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
|
|
|
|
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, Subset
|
|
from tqdm.auto import tqdm
|
|
|
|
from dataset import _ExpoBaseDataset
|
|
from evaluate_auc_v2 import (
|
|
build_metadata_for_merge,
|
|
build_model_from_dataset,
|
|
get_auc_delong_var,
|
|
load_checkpoint_state_dict,
|
|
load_json_config,
|
|
load_model_state,
|
|
parse_float_list,
|
|
parse_int_list,
|
|
project_distribution_chunk,
|
|
resolve_dist_mode_for_checkpoint,
|
|
select_disease_tokens,
|
|
validate_dataset_metadata,
|
|
_score_to_probability,
|
|
)
|
|
from readouts import build_readout
|
|
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX, DAYS_PER_YEAR
|
|
|
|
|
|
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
|
|
|
|
|
def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any:
|
|
value = getattr(args, name, None)
|
|
if value is not None:
|
|
return value
|
|
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,
|
|
data_prefix: str,
|
|
labels_file: str,
|
|
model_target_mode: str,
|
|
extra_info_types: Iterable[int] | None = None,
|
|
) -> None:
|
|
super().__init__(
|
|
data_prefix=data_prefix,
|
|
labels_file=labels_file,
|
|
no_event_interval_years=5.0,
|
|
include_no_event_in_uts_target=False,
|
|
extra_info_types=extra_info_types,
|
|
)
|
|
self.model_target_mode = str(model_target_mode).lower()
|
|
if self.model_target_mode not in {"next_token", "all_future"}:
|
|
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
|
|
|
|
self.records: List[Dict[str, Any]] = []
|
|
self.first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
|
|
|
unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True)
|
|
ends = np.concatenate([starts[1:], [len(self.event_data)]])
|
|
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
|
|
|
for eid_raw, start, end in zip(unique_eids, starts, ends):
|
|
eid = int(eid_raw)
|
|
rows = self.event_data[start:end]
|
|
checkup_rows = rows[rows[:, 2].astype(np.int64) == CHECKUP_IDX]
|
|
if len(checkup_rows) == 0:
|
|
continue
|
|
|
|
features = self._split_features(eid)
|
|
if features is None:
|
|
continue
|
|
|
|
doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32)))
|
|
doa_years = np.float32(doa_days / DAYS_PER_YEAR)
|
|
|
|
disease_rows = rows[rows[:, 2].astype(np.int64) != CHECKUP_IDX]
|
|
disease_times = disease_rows[:, 1].astype(np.float32) / DAYS_PER_YEAR
|
|
disease_labels_raw = disease_rows[:, 2].astype(np.int64)
|
|
disease_labels = np.where(
|
|
disease_labels_raw >= NO_EVENT_IDX,
|
|
disease_labels_raw + 1,
|
|
disease_labels_raw,
|
|
).astype(np.int64)
|
|
order = np.lexsort((disease_labels, disease_times))
|
|
disease_times = disease_times[order].astype(np.float32)
|
|
disease_labels = disease_labels[order].astype(np.int64)
|
|
|
|
patient_id = len(self.records)
|
|
for token in np.unique(disease_labels).tolist():
|
|
token = int(token)
|
|
if token in SPECIAL_TOKENS:
|
|
continue
|
|
hit = np.where(disease_labels == token)[0]
|
|
if hit.size:
|
|
first_lists.setdefault(token, []).append(
|
|
(patient_id, float(disease_times[int(hit[0])]))
|
|
)
|
|
|
|
hist = disease_times <= doa_years
|
|
hist_events = disease_labels[hist]
|
|
hist_times = disease_times[hist]
|
|
|
|
if self.model_target_mode == "next_token":
|
|
at_doa = np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6)
|
|
if hist_events.size == 0 or not np.any(at_doa):
|
|
event_seq = np.concatenate([
|
|
hist_events,
|
|
np.array([NO_EVENT_IDX], dtype=np.int64),
|
|
])
|
|
time_seq = np.concatenate([
|
|
hist_times,
|
|
np.array([doa_years], dtype=np.float32),
|
|
])
|
|
else:
|
|
event_seq = hist_events
|
|
time_seq = hist_times
|
|
readout_pos = int(len(event_seq) - 1)
|
|
else:
|
|
event_seq = hist_events
|
|
time_seq = hist_times
|
|
readout_pos = -1
|
|
|
|
self.records.append(
|
|
{
|
|
"patient_id": patient_id,
|
|
"eid": eid,
|
|
"doa": doa_years,
|
|
"event_seq": event_seq.astype(np.int64),
|
|
"time_seq": time_seq.astype(np.float32),
|
|
"readout_pos": readout_pos,
|
|
"full_events": disease_labels,
|
|
"full_times": disease_times,
|
|
"sex": int(features["sex"]),
|
|
"other_type": features["other_type"],
|
|
"other_value": features["other_value"],
|
|
"other_value_kind": features["other_value_kind"],
|
|
"other_time": features["other_time"],
|
|
}
|
|
)
|
|
|
|
for token, pairs in first_lists.items():
|
|
self.first_occurrence_by_token[int(token)] = (
|
|
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
|
np.asarray([t for _, t in pairs], dtype=np.float32),
|
|
)
|
|
|
|
if not self.records:
|
|
raise RuntimeError("No DOA records were built from checkup events.")
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.records)
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
|
s = self.records[idx]
|
|
return {
|
|
"event_seq": torch.from_numpy(s["event_seq"]).long(),
|
|
"time_seq": torch.from_numpy(s["time_seq"]).float(),
|
|
"readout_pos": torch.tensor(s["readout_pos"], dtype=torch.long),
|
|
"t_query": torch.tensor(float(s["doa"]), dtype=torch.float32),
|
|
"patient_id": torch.tensor(s["patient_id"], dtype=torch.long),
|
|
"sex": torch.tensor(s["sex"], dtype=torch.long),
|
|
"other_type": torch.from_numpy(s["other_type"]).long(),
|
|
"other_value": torch.from_numpy(s["other_value"]).float(),
|
|
"other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
|
|
"other_time": torch.from_numpy(s["other_time"]).float(),
|
|
}
|
|
|
|
|
|
def collate_doa_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
|
event_seq = pad_sequence(
|
|
[x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
|
|
)
|
|
time_seq = pad_sequence(
|
|
[x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
other_type = pad_sequence(
|
|
[x["other_type"] for x in batch], batch_first=True, padding_value=0
|
|
)
|
|
other_value = pad_sequence(
|
|
[x["other_value"] for x in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
other_value_kind = pad_sequence(
|
|
[x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
|
|
)
|
|
other_time = pad_sequence(
|
|
[x["other_time"] for x in batch], batch_first=True, padding_value=0.0
|
|
)
|
|
|
|
readout_mask = torch.zeros_like(event_seq, dtype=torch.bool)
|
|
readout_pos = torch.stack([x["readout_pos"] for x in batch])
|
|
for i, pos in enumerate(readout_pos.tolist()):
|
|
if pos >= 0:
|
|
readout_mask[i, int(pos)] = True
|
|
|
|
return {
|
|
"event_seq": event_seq,
|
|
"time_seq": time_seq,
|
|
"padding_mask": event_seq > PAD_IDX,
|
|
"readout_mask": readout_mask,
|
|
"readout_pos": readout_pos,
|
|
"t_query": torch.stack([x["t_query"] for x in batch]),
|
|
"patient_id": torch.stack([x["patient_id"] for x in batch]),
|
|
"sex": torch.stack([x["sex"] for x in batch]),
|
|
"other_type": other_type,
|
|
"other_value": other_value,
|
|
"other_value_kind": other_value_kind,
|
|
"other_time": other_time,
|
|
}
|
|
|
|
|
|
@torch.inference_mode()
|
|
def infer_doa_hidden(
|
|
model,
|
|
loader: DataLoader,
|
|
device: torch.device,
|
|
model_target_mode: str,
|
|
readout_name: str,
|
|
readout_reduce: str,
|
|
use_amp: bool,
|
|
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
|
|
model_target_mode = str(model_target_mode).lower()
|
|
readout = None
|
|
if model_target_mode == "next_token":
|
|
if readout_name == "same_time_group_end":
|
|
readout = build_readout("same_time_group_end", reduce=readout_reduce).to(device)
|
|
else:
|
|
readout = build_readout(readout_name).to(device)
|
|
readout.eval()
|
|
|
|
hidden_parts: List[np.ndarray] = []
|
|
patient_parts: List[np.ndarray] = []
|
|
sex_parts: List[np.ndarray] = []
|
|
autocast_enabled = bool(use_amp and device.type == "cuda")
|
|
|
|
for batch in tqdm(loader, desc="DOA inference", leave=False, dynamic_ncols=True):
|
|
batch_dev = {
|
|
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
|
for k, v in batch.items()
|
|
}
|
|
amp_context = (
|
|
torch.autocast(device_type=device.type, dtype=torch.float16)
|
|
if autocast_enabled
|
|
else contextlib.nullcontext()
|
|
)
|
|
with amp_context:
|
|
if model_target_mode == "all_future":
|
|
hidden = model(
|
|
event_seq=batch_dev["event_seq"],
|
|
time_seq=batch_dev["time_seq"],
|
|
sex=batch_dev["sex"],
|
|
padding_mask=batch_dev["padding_mask"],
|
|
t_query=batch_dev["t_query"],
|
|
other_type=batch_dev["other_type"],
|
|
other_value=batch_dev["other_value"],
|
|
other_value_kind=batch_dev["other_value_kind"],
|
|
other_time=batch_dev["other_time"],
|
|
target_mode="all_future",
|
|
)
|
|
else:
|
|
hidden_raw = model(
|
|
event_seq=batch_dev["event_seq"],
|
|
time_seq=batch_dev["time_seq"],
|
|
sex=batch_dev["sex"],
|
|
padding_mask=batch_dev["padding_mask"],
|
|
other_type=batch_dev["other_type"],
|
|
other_value=batch_dev["other_value"],
|
|
other_value_kind=batch_dev["other_value_kind"],
|
|
other_time=batch_dev["other_time"],
|
|
target_mode="next_token",
|
|
)
|
|
ro = readout(
|
|
hidden=hidden_raw,
|
|
time_seq=batch_dev["time_seq"],
|
|
padding_mask=batch_dev["padding_mask"],
|
|
readout_mask=batch_dev["readout_mask"],
|
|
)
|
|
if ro.hidden.dim() == 2:
|
|
hidden = ro.hidden
|
|
else:
|
|
hidden = ro.hidden[batch_dev["readout_mask"]]
|
|
|
|
hidden_parts.append(hidden.detach().float().cpu().numpy().astype(np.float32, copy=False))
|
|
patient_parts.append(batch["patient_id"].cpu().numpy().astype(np.int32, copy=False))
|
|
sex_parts.append(batch["sex"].cpu().numpy().astype(np.int8, copy=False))
|
|
|
|
return (
|
|
np.concatenate(hidden_parts, axis=0),
|
|
{
|
|
"patient_id": np.concatenate(patient_parts, axis=0),
|
|
"sex": np.concatenate(sex_parts, axis=0),
|
|
},
|
|
)
|
|
|
|
|
|
def first_time_array(
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
token: int,
|
|
patient_count: int,
|
|
) -> np.ndarray:
|
|
out = np.full(patient_count, np.inf, dtype=np.float32)
|
|
pairs = first_occurrence_by_token.get(int(token))
|
|
if pairs is not None:
|
|
p, t = pairs
|
|
out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
|
|
return out
|
|
|
|
|
|
_DOA_WORKER: Dict[str, Any] = {}
|
|
|
|
|
|
def _init_doa_worker(
|
|
disease_ids: np.ndarray,
|
|
logits_all: np.ndarray,
|
|
rho_all: Optional[np.ndarray],
|
|
row_patient_id: np.ndarray,
|
|
row_sex: np.ndarray,
|
|
row_doa: np.ndarray,
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
patient_count: int,
|
|
horizons: np.ndarray,
|
|
min_cases: int,
|
|
dist_mode: str,
|
|
score_mode: str,
|
|
death_idx: int,
|
|
) -> None:
|
|
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
|
os.environ.setdefault("MKL_NUM_THREADS", "1")
|
|
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
|
|
os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
|
|
_DOA_WORKER.clear()
|
|
_DOA_WORKER.update(
|
|
{
|
|
"disease_ids": np.asarray(disease_ids, dtype=np.int64),
|
|
"logits_all": np.asarray(logits_all, dtype=np.float32),
|
|
"rho_all": None if rho_all is None else np.asarray(rho_all, dtype=np.float32),
|
|
"row_patient_id": np.asarray(row_patient_id, dtype=np.int32),
|
|
"row_sex": np.asarray(row_sex, dtype=np.int8),
|
|
"row_doa": np.asarray(row_doa, dtype=np.float32),
|
|
"first_occurrence_by_token": first_occurrence_by_token,
|
|
"patient_count": int(patient_count),
|
|
"horizons": np.asarray(horizons, dtype=np.float32),
|
|
"min_cases": int(min_cases),
|
|
"dist_mode": str(dist_mode).lower(),
|
|
"score_mode": str(score_mode).lower(),
|
|
"death_idx": int(death_idx),
|
|
"first_time_cache": {},
|
|
}
|
|
)
|
|
|
|
|
|
def _doa_first_time_by_patient(token: int) -> np.ndarray:
|
|
cache = _DOA_WORKER["first_time_cache"]
|
|
if int(token) in cache:
|
|
return cache[int(token)]
|
|
|
|
out = np.full(int(_DOA_WORKER["patient_count"]), np.inf, dtype=np.float32)
|
|
pairs = _DOA_WORKER["first_occurrence_by_token"].get(int(token))
|
|
if pairs is not None:
|
|
p, t = pairs
|
|
out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32)
|
|
cache[int(token)] = out
|
|
return out
|
|
|
|
|
|
def _eval_doa_token(task: Tuple[int, int]) -> List[Dict[str, Any]]:
|
|
col, token = task
|
|
col = int(col)
|
|
token = int(token)
|
|
|
|
patient_ids = _DOA_WORKER["row_patient_id"]
|
|
sex = _DOA_WORKER["row_sex"]
|
|
doa = _DOA_WORKER["row_doa"]
|
|
logits = _DOA_WORKER["logits_all"][:, col]
|
|
rho_all = _DOA_WORKER["rho_all"]
|
|
rho = None if rho_all is None else rho_all[:, col]
|
|
first_time = _doa_first_time_by_patient(token)[patient_ids]
|
|
never = np.isinf(first_time)
|
|
incident_after_doa = first_time > doa
|
|
|
|
rows: List[Dict[str, Any]] = []
|
|
for horizon in _DOA_WORKER["horizons"].tolist():
|
|
horizon = float(horizon)
|
|
case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon))
|
|
control_mask = never
|
|
if int(case_mask.sum()) < int(_DOA_WORKER["min_cases"]) or int(control_mask.sum()) < int(_DOA_WORKER["min_cases"]):
|
|
continue
|
|
|
|
scores = _score_to_probability(
|
|
logits=logits,
|
|
rho=rho,
|
|
score_mode=_DOA_WORKER["score_mode"],
|
|
horizon=horizon,
|
|
dist_mode=_DOA_WORKER["dist_mode"],
|
|
token=token,
|
|
death_idx=int(_DOA_WORKER["death_idx"]),
|
|
)
|
|
|
|
for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]:
|
|
sex_mask = np.ones_like(case_mask, dtype=bool) if sex_value == -1 else sex == sex_value
|
|
cm = case_mask & sex_mask
|
|
nm = control_mask & sex_mask
|
|
if int(cm.sum()) < int(_DOA_WORKER["min_cases"]) or int(nm.sum()) < int(_DOA_WORKER["min_cases"]):
|
|
continue
|
|
auc, var = get_auc_delong_var(scores[cm], scores[nm])
|
|
rows.append(
|
|
{
|
|
"token": token,
|
|
"horizon": horizon,
|
|
"sex": sex_name,
|
|
"n_case": int(cm.sum()),
|
|
"n_control": int(nm.sum()),
|
|
"auc": auc,
|
|
"auc_var": var,
|
|
"auc_se": float(np.sqrt(max(var, 0.0))) if np.isfinite(var) else np.nan,
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
def _doa_task_block(tasks: Sequence[Tuple[int, int]]) -> List[Dict[str, Any]]:
|
|
rows: List[Dict[str, Any]] = []
|
|
for task in tasks:
|
|
rows.extend(_eval_doa_token(task))
|
|
return rows
|
|
|
|
|
|
def _split_tasks(tasks: Sequence[Tuple[int, int]], chunk_size: int) -> List[List[Tuple[int, int]]]:
|
|
if not tasks:
|
|
return []
|
|
if chunk_size <= 0:
|
|
chunk_size = max(1, int(np.ceil(len(tasks) / 8)))
|
|
return [list(tasks[i:i + chunk_size]) for i in range(0, len(tasks), chunk_size)]
|
|
|
|
|
|
def evaluate_doa_auc_chunk(
|
|
dataset: DOAStatusDataset,
|
|
hidden_all: np.ndarray,
|
|
row_arrays: Dict[str, np.ndarray],
|
|
model,
|
|
disease_ids: Sequence[int],
|
|
horizons: np.ndarray,
|
|
dist_mode: str,
|
|
score_mode: str,
|
|
min_cases: int,
|
|
device: torch.device,
|
|
logit_batch_size: int,
|
|
use_amp: bool,
|
|
num_workers_auc: int,
|
|
auc_task_chunk_size: int,
|
|
) -> pd.DataFrame:
|
|
logits_all, rho_all = project_distribution_chunk(
|
|
model=model,
|
|
hidden_all=hidden_all,
|
|
disease_ids=disease_ids,
|
|
dist_mode=dist_mode,
|
|
device=device,
|
|
logit_batch_size=logit_batch_size,
|
|
use_amp=use_amp,
|
|
)
|
|
patient_ids = row_arrays["patient_id"].astype(np.int32)
|
|
doa = np.asarray([r["doa"] for r in dataset.records], dtype=np.float32)[patient_ids]
|
|
patient_count = len(dataset.records)
|
|
death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
|
|
disease_ids_arr = np.asarray([int(x) for x in disease_ids], dtype=np.int64)
|
|
tasks = [(j, int(token)) for j, token in enumerate(disease_ids_arr.tolist())]
|
|
|
|
init_args = (
|
|
disease_ids_arr,
|
|
logits_all,
|
|
rho_all,
|
|
patient_ids,
|
|
row_arrays["sex"].astype(np.int8),
|
|
doa,
|
|
dataset.first_occurrence_by_token,
|
|
patient_count,
|
|
horizons,
|
|
min_cases,
|
|
dist_mode,
|
|
score_mode,
|
|
death_idx,
|
|
)
|
|
|
|
if int(num_workers_auc) <= 1 or len(tasks) <= 1:
|
|
_init_doa_worker(*init_args)
|
|
rows = _doa_task_block(tasks)
|
|
return pd.DataFrame(rows)
|
|
|
|
rows: List[Dict[str, Any]] = []
|
|
task_blocks = _split_tasks(tasks, int(auc_task_chunk_size))
|
|
with ProcessPoolExecutor(
|
|
max_workers=int(num_workers_auc),
|
|
initializer=_init_doa_worker,
|
|
initargs=init_args,
|
|
) as pool:
|
|
futures = [pool.submit(_doa_task_block, block) for block in task_blocks]
|
|
for fut in tqdm(futures, desc="DOA AUC workers", leave=False, dynamic_ncols=True):
|
|
rows.extend(fut.result())
|
|
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("--num_workers_auc", type=int, default=None)
|
|
parser.add_argument("--auc_task_chunk_size", 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)
|
|
parser.add_argument("--min_cases", type=int, default=None)
|
|
parser.add_argument("--labels_meta_path", type=str, default=None)
|
|
parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None)
|
|
args = parser.parse_args()
|
|
|
|
run_path = Path(args.run_path)
|
|
cfg = load_json_config(run_path / "train_config.json")
|
|
ckpt_path = run_path / "best_model.pt"
|
|
if not ckpt_path.exists():
|
|
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
|
|
|
output_path = Path(args.output_path or run_path)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
|
if model_target_mode not in {"next_token", "all_future"}:
|
|
raise ValueError(f"Unsupported model_target_mode={model_target_mode!r}")
|
|
|
|
labels_meta_path = cfg_get(args, cfg, "labels_meta_path", None)
|
|
if labels_meta_path is None:
|
|
labels_meta_path = cfg.get("labels_meta_path", "delphi_labels_chapters_colours_icd.csv")
|
|
labels_meta = pd.read_csv(labels_meta_path) if labels_meta_path and Path(labels_meta_path).exists() else None
|
|
|
|
dataset = DOAStatusDataset(
|
|
data_prefix=cfg.get("data_prefix", "ukb"),
|
|
labels_file=cfg.get("labels_file", "labels.csv"),
|
|
model_target_mode=model_target_mode,
|
|
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(
|
|
dataset=dataset,
|
|
labels_meta=labels_meta,
|
|
requested_tokens=disease_requested,
|
|
filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
|
|
first_occurrence_by_token=eval_first_occurrence,
|
|
)
|
|
if not disease_ids:
|
|
raise RuntimeError("No disease tokens selected after filtering.")
|
|
|
|
horizons = np.asarray(
|
|
parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0],
|
|
dtype=np.float32,
|
|
)
|
|
score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower()
|
|
min_cases = int(cfg_get(args, cfg, "min_cases", 2))
|
|
|
|
state_dict = load_checkpoint_state_dict(ckpt_path, map_location="cpu")
|
|
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
|
|
cfg_model = dict(cfg)
|
|
cfg_model["dist_mode"] = dist_mode
|
|
|
|
device = torch.device(cfg.get("device", "cuda") if torch.cuda.is_available() else "cpu")
|
|
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
|
load_model_state(model, state_dict)
|
|
model.eval()
|
|
|
|
if model_target_mode == "next_token" and (
|
|
model.token_embedding.num_embeddings <= NO_EVENT_IDX
|
|
or model.risk_head.out_features <= NO_EVENT_IDX
|
|
):
|
|
raise RuntimeError("Next-token DOA evaluation requires <NO_EVENT> in the model vocabulary.")
|
|
|
|
eval_dataset = Subset(dataset, eval_indices)
|
|
loader = DataLoader(
|
|
eval_dataset,
|
|
batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
|
|
shuffle=False,
|
|
collate_fn=collate_doa_fn,
|
|
num_workers=int(cfg_get(args, cfg, "num_workers", 4)),
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0,
|
|
prefetch_factor=2 if int(cfg_get(args, cfg, "num_workers", 4)) > 0 else None,
|
|
)
|
|
|
|
target_mode = cfg.get("target_mode", "uts")
|
|
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"))
|
|
|
|
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}")
|
|
print(f"Horizons: {horizons.tolist()}")
|
|
print(f"Disease tokens: {len(disease_ids)}")
|
|
|
|
hidden_all, row_arrays = infer_doa_hidden(
|
|
model=model,
|
|
loader=loader,
|
|
device=device,
|
|
model_target_mode=model_target_mode,
|
|
readout_name=readout_name,
|
|
readout_reduce=readout_reduce,
|
|
use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
|
|
)
|
|
chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 256))
|
|
num_workers_auc = int(cfg_get(args, cfg, "num_workers_auc", max(1, (os.cpu_count() or 2) - 1)))
|
|
auc_task_chunk_size = int(cfg_get(args, cfg, "auc_task_chunk_size", 0))
|
|
print(f"Disease chunk size: {chunk_size}")
|
|
print(f"AUC workers: {num_workers_auc}")
|
|
result_parts = []
|
|
for disease_chunk in tqdm(
|
|
iter_chunks(disease_ids, chunk_size),
|
|
desc="Disease chunks",
|
|
leave=True,
|
|
dynamic_ncols=True,
|
|
):
|
|
result_parts.append(
|
|
evaluate_doa_auc_chunk(
|
|
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)),
|
|
num_workers_auc=num_workers_auc,
|
|
auc_task_chunk_size=auc_task_chunk_size,
|
|
)
|
|
)
|
|
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")
|
|
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(
|
|
auc_mean=("auc", "mean"),
|
|
n_case=("n_case", "sum"),
|
|
n_control=("n_control", "sum"),
|
|
)
|
|
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}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|