"""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 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 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 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) 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 def evaluate_doa_auc( 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, ) -> 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) sex = row_arrays["sex"].astype(np.int8) 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)) rows: List[Dict[str, Any]] = [] for col, token in enumerate([int(x) for x in disease_ids]): first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids] never = np.isinf(first_time) incident_after_doa = first_time > doa for horizon in horizons.tolist(): horizon = float(horizon) case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon)) control_mask = never if int(case_mask.sum()) < min_cases or int(control_mask.sum()) < min_cases: continue rho_col = None if rho_all is None else rho_all[:, col] scores = _score_to_probability( logits=logits_all[:, col], rho=rho_col, score_mode=score_mode, horizon=horizon, dist_mode=dist_mode, token=token, death_idx=death_idx, ) for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]: if sex_value == -1: sex_mask = np.ones_like(case_mask, dtype=bool) else: sex_mask = sex == sex_value cm = case_mask & sex_mask nm = control_mask & sex_mask if int(cm.sum()) < min_cases or int(nm.sum()) < 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 pd.DataFrame(rows) 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("--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("--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) 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=dataset.first_occurrence_by_token, ) 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 in the model vocabulary.") loader = DataLoader( 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")) print(f"DOA records: {len(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)), ) 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)), ) 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.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.to_csv(output_path / "doa_auc_summary.csv", index=False) print(f"Wrote {out_file}") if __name__ == "__main__": main()