diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..eeeec40 --- /dev/null +++ b/LICENSE @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2026, Jiarui Li + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT INCLUDING NEGLIGENCE OR OTHERWISE ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/evaluate_auc.py b/evaluate_auc.py new file mode 100644 index 0000000..d96e70c --- /dev/null +++ b/evaluate_auc.py @@ -0,0 +1,1327 @@ +""" +Evaluate disease-specific AUC for DeepHealth models. + +This script follows the logic of the Delphi evaluation script supplied by the user: + 1. choose disease tokens of interest from dataset vocabulary; + 2. precompute, for each patient/target occurrence, the latest prediction token + at least `offset` years before the target time; + 3. run model inference by disease chunks to avoid materializing all logits; + 4. compute AUC separately by sex and age bracket; + 5. aggregate age brackets with DeLong variance. + +Efficiency notes: + - transformer/readout inference is executed once and cached; + - disease chunks reuse the cached hidden states and only recompute selected risk-head columns; + - AUC work is parallelized on CPU across disease task blocks using process workers; + - per-sex data are compacted into an event-level table before multiprocessing; + - large per-sex arrays are installed once per worker with fork-style globals on Linux, + avoiding repeated pickling of arrays for every disease. + +Run from the DeepHealth code directory containing dataset.py, models.py, +readouts.py, and train.py-compatible checkpoints/configs. +""" +from __future__ import annotations + +import argparse +import contextlib +import json +import math +import multiprocessing as mp +import os +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader, Subset +from tqdm.auto import tqdm + +from dataset import HealthDataset, collate_fn +from models import DeepHealth +from readouts import build_readout +from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX + + +# --------------------------------------------------------------------------- +# DeLong AUC utilities, adapted from the supplied Delphi evaluation script +# --------------------------------------------------------------------------- + +def compute_midrank(x: np.ndarray) -> np.ndarray: + """Compute midranks for DeLong variance.""" + x = np.asarray(x, dtype=np.float64) + order = np.argsort(x) + sorted_x = x[order] + n = len(x) + ranks = np.zeros(n, dtype=np.float64) + i = 0 + while i < n: + j = i + while j < n and sorted_x[j] == sorted_x[i]: + j += 1 + ranks[i:j] = 0.5 * (i + j - 1) + i = j + out = np.empty(n, dtype=np.float64) + out[order] = ranks + 1.0 + return out + + +def fast_delong(predictions_sorted_transposed: np.ndarray, label_1_count: int) -> Tuple[np.ndarray, np.ndarray]: + """ + Fast DeLong covariance for ROC AUC. + + This evaluation uses one classifier at a time, so k is normally 1. + In that case np.cov(v01) and np.cov(v10) must be scalar variances. + Do NOT expand them to (m,m)/(n,n) covariance matrices; that is the + source of the broadcast error: + shapes (n_case,n_case) and (n_control,n_control) + """ + predictions_sorted_transposed = np.asarray( + predictions_sorted_transposed, dtype=np.float64) + m = int(label_1_count) + n = int(predictions_sorted_transposed.shape[1] - m) + if m <= 0 or n <= 0: + return np.array([np.nan], dtype=np.float64), np.array([[np.nan]], dtype=np.float64) + + positive_examples = predictions_sorted_transposed[:, :m] + negative_examples = predictions_sorted_transposed[:, m:] + k = int(predictions_sorted_transposed.shape[0]) + + tx = np.empty((k, m), dtype=np.float64) + ty = np.empty((k, n), dtype=np.float64) + tz = np.empty((k, m + n), dtype=np.float64) + for r in range(k): + tx[r] = compute_midrank(positive_examples[r]) + ty[r] = compute_midrank(negative_examples[r]) + tz[r] = compute_midrank(predictions_sorted_transposed[r]) + + aucs = tz[:, :m].sum(axis=1) / m / n - float(m + 1.0) / 2.0 / n + + # DeLong structural components: + # v01: classifier x cases + # v10: classifier x controls + v01 = (tz[:, :m] - tx) / n + v10 = 1.0 - (tz[:, m:] - ty) / m + + if k == 1: + # np.cov on a single row is easy to misuse. The correct covariance + # for one classifier is simply the scalar sample variance over cases + # plus the scalar sample variance over controls. + sx = np.var(v01[0], ddof=1) if m > 1 else 0.0 + sy = np.var(v10[0], ddof=1) if n > 1 else 0.0 + delong_cov = np.array([[sx / m + sy / n]], dtype=np.float64) + else: + # Multiple-classifier general case. rowvar=True: rows are classifiers. + sx = np.cov(v01, rowvar=True) if m > 1 else np.zeros( + (k, k), dtype=np.float64) + sy = np.cov(v10, rowvar=True) if n > 1 else np.zeros( + (k, k), dtype=np.float64) + delong_cov = np.atleast_2d(sx) / m + np.atleast_2d(sy) / n + + return aucs, delong_cov + + +def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> Tuple[float, float]: + """Return AUC and DeLong variance for controls/class-0 and cases/class-1.""" + control_scores = np.asarray(control_scores, dtype=np.float64) + case_scores = np.asarray(case_scores, dtype=np.float64) + if len(control_scores) == 0 or len(case_scores) == 0: + return np.nan, np.nan + + ground_truth = np.array([1] * len(case_scores) + + [0] * len(control_scores), dtype=np.int8) + predictions = np.concatenate( + [case_scores, control_scores]).astype(np.float64, copy=False) + order = (-ground_truth).argsort() + label_1_count = int(ground_truth.sum()) + + aucs, delong_cov = fast_delong( + predictions[np.newaxis, order], label_1_count) + var = float(np.asarray(delong_cov).reshape(-1)[0]) + return float(aucs[0]), var + + +# --------------------------------------------------------------------------- +# Disease selection +# --------------------------------------------------------------------------- + +SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} + + +def _build_first_occurrence_maps( + dataset: HealthDataset, + subset_indices: np.ndarray, +) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]: + patient_count = len(subset_indices) + followup_end = np.full(patient_count, -np.inf, dtype=np.float32) + death_time = np.full(patient_count, np.inf, dtype=np.float32) + sex = np.full(patient_count, -1, dtype=np.int8) + + first_lists: Dict[int, List[Tuple[int, float]]] = {} + + for patient_id, dataset_index in enumerate(subset_indices.tolist()): + s = dataset.samples[int(dataset_index)] + + seq_event = np.asarray(s["event_seq"], dtype=np.int64) + seq_time = np.asarray(s["time_seq"], dtype=np.float32) + tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64) + tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32) + + if seq_event.size == 0 or tgt_event.size == 0: + continue + + full_event = np.concatenate([seq_event, tgt_event[-1:]]) + full_time = np.concatenate([seq_time, tgt_time[-1:]]) + + sex[patient_id] = int(s["sex"]) + followup_end[patient_id] = np.max(full_time).astype(np.float32) + + uniq_tokens, first_idx = np.unique(full_event, return_index=True) + + for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()): + token = int(token) + event_time = float(full_time[int(idx)]) + if token not in first_lists: + first_lists[token] = [] + first_lists[token].append((patient_id, event_time)) + + packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {} + for token, pairs in first_lists.items(): + if not pairs: + continue + packed[int(token)] = ( + np.asarray([p for p, _ in pairs], dtype=np.int32), + np.asarray([t for _, t in pairs], dtype=np.float32), + ) + + return packed, followup_end, death_time, sex + + +def _get_death_token_ids(dataset: HealthDataset) -> List[int]: + ids: List[int] = [] + + exact_codes = {"death", "", "dth", "deceased", "mortality"} + for token, code in dataset.label_id_to_code.items(): + token = int(token) + if token in SPECIAL_TOKENS: + continue + text = str(code).strip().lower() + if text in exact_codes or ("death" in text) or ("mortality" in text): + ids.append(token) + + death_ids = sorted(set(int(x) + for x in ids if int(x) not in SPECIAL_TOKENS)) + print(f"[INFO] death token ids: {death_ids}") + return death_ids + + +def select_disease_tokens( + dataset: HealthDataset, + requested_tokens: Optional[Sequence[int]], + filter_min_total: int, + first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], +) -> List[int]: + base = [ + int(token) + for token, code in dataset.label_id_to_code.items() + if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<") + ] + base_set = set(base) + print(f"[INFO] Valid disease tokens in current vocabulary: {len(base)}") + + if requested_tokens is not None: + selected = sorted( + set(int(x) for x in requested_tokens if int(x) in base_set and int(x) not in SPECIAL_TOKENS)) + print( + "[INFO] Requested disease tokens provided: " + f"input={len(list(requested_tokens))}, selected={len(selected)} (vocab/SPECIAL filtered).") + print(f"[INFO] Final disease_ids count: {len(selected)}") + return selected + + disease_ids = sorted(base) + if int(filter_min_total) <= 0: + print( + f"[INFO] filter_min_total={int(filter_min_total)} <= 0; keeping all {len(disease_ids)} disease tokens.") + print(f"[INFO] Final disease_ids count: {len(disease_ids)}") + return disease_ids + + print( + f"[INFO] Applying filter_min_total={int(filter_min_total)}: before={len(disease_ids)} tokens.") + print( + "[INFO] Using split first-occurrence patient counts for filter_min_total.") + split_counts = {} + for token, pairs in first_occurrence_by_token.items(): + token = int(token) + if token not in base_set: + continue + split_counts[token] = len(np.unique(pairs[0])) + filtered = [token for token in disease_ids if int( + split_counts.get(token, 0)) > int(filter_min_total)] + print( + "[INFO] First-occurrence count filtering complete: " + f"after={len(filtered)} tokens.") + print(f"[INFO] Final disease_ids count: {len(filtered)}") + return filtered + + +# --------------------------------------------------------------------------- +# Dataset/split/model helpers +# --------------------------------------------------------------------------- + +def load_json_config(path: Optional[str]) -> Dict[str, Any]: + if path is None: + return {} + p = Path(path) + if not p.exists(): + return {} + with p.open("r", encoding="utf-8") as f: + return json.load(f) + + +def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any], name: str, default: Any) -> Any: + """Get a value from CLI args first, then train_config.json, then default. + + This helper intentionally accepts either an argparse.Namespace or a dict. + The earlier version passed cfg as both args and cfg, then tried to access + args.eval_split, which fails because dict has no attributes. + """ + val = None + if args is not None: + if isinstance(args, dict): + val = args.get(name, None) + else: + val = getattr(args, name, None) + if val is not None: + return val + 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 = train_ratio + val_ratio + 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}") + rng = np.random.RandomState(seed) + idx = rng.permutation(n) + n_train = int(n * train_ratio) + n_val = int(n * val_ratio) + return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:] + + +def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth: + return DeepHealth( + vocab_size=dataset.vocab_size, + n_embd=int(cfg_get(args, cfg, "n_embd", 120)), + n_head=int(cfg_get(args, cfg, "n_head", 10)), + n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)), + n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)), + n_types=dataset.n_types, + n_cont_types=dataset.n_cont_types, + n_categories=dataset.n_categories, + cont_type_ids=dataset.cont_type_ids, + n_bins=int(cfg_get(args, cfg, "n_bins", 16)), + target_mode="next_token", + time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), + dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), + dropout=float(cfg_get(args, cfg, "dropout", 0.0)), + ) + + +def _extract_state_dict(ckpt: Any) -> Dict[str, Any]: + if isinstance(ckpt, dict) and "model" in ckpt: + return ckpt["model"] + elif isinstance(ckpt, dict) and "state_dict" in ckpt: + return ckpt["state_dict"] + return ckpt + + +def load_checkpoint_state_dict(checkpoint_path: str, map_location: str | torch.device = "cpu") -> Dict[str, Any]: + ckpt = torch.load(checkpoint_path, map_location=map_location) + state = _extract_state_dict(ckpt) + if not isinstance(state, dict): + raise TypeError( + f"Unsupported checkpoint payload type: {type(state)}") + return state + + +def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str: + mode = str(cfg_dist_mode).lower() + has_rho_head = any(str(k).startswith("rho_head.") + for k in state_dict.keys()) + + if has_rho_head and mode != "weibull": + print( + "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.") + return "weibull" + if (not has_rho_head) and mode == "weibull": + print( + "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.") + return "exponential" + return mode + + +def load_model_state( + model: torch.nn.Module, + checkpoint_path: str, + device: torch.device, + state_dict: Optional[Dict[str, Any]] = None, +) -> None: + state = state_dict if state_dict is not None else load_checkpoint_state_dict( + checkpoint_path, map_location=device) + + missing, unexpected = model.load_state_dict(state, strict=False) + if missing or unexpected: + print( + f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}") + + +def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, 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() + dataset_subset_size = cfg_get(args, cfg, "dataset_subset_size", None) + + 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, + "valid": val_idx, + "validation": val_idx, + "test": test_idx, + "all": np.arange(len(dataset)), + } + if eval_split not in split_map: + raise ValueError( + f"eval_split must be one of {sorted(split_map)}, got {eval_split!r}") + + indices = split_map[eval_split] + if dataset_subset_size is not None and int(dataset_subset_size) > 0: + indices = indices[: int(dataset_subset_size)] + return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64) + + +# --------------------------------------------------------------------------- +# Batched inference + cached hidden states +# --------------------------------------------------------------------------- + + +def _numpy_hidden_dtype(name: str) -> np.dtype: + name = str(name).lower() + if name in {"float16", "fp16", "half"}: + return np.float16 + if name in {"bfloat16", "bf16"}: + # NumPy has limited bfloat16 support; store as fp16 for compact CPU cache. + return np.float16 + if name in {"float32", "fp32", "single"}: + return np.float32 + raise ValueError( + f"hidden_cache_dtype must be float16 or float32, got {name!r}") + + +@torch.inference_mode() +def infer_readout_hidden( + model: DeepHealth, + loader: DataLoader, + device: torch.device, + attn_mask_mode: str, + readout_name: str, + readout_reduce: str, + use_amp: bool, + hidden_cache_dtype: str = "float16", +) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: + """ + Run the expensive transformer/readout path exactly once and cache hidden states. + + The older implementation re-ran the whole model once per disease chunk even + though only the final tied risk-head columns changed. That is usually the + dominant bottleneck. This function computes readout hidden states once; later + chunks only perform a cheap selected-vocabulary linear projection. + """ + 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] = [] + arrays: Dict[str, List[np.ndarray]] = { + "event_seq": [], + "time_seq": [], + "target_event_seq": [], + "target_time_seq": [], + "padding_mask": [], + "readout_mask": [], + "sex": [], + } + max_len = 0 + out_dtype = _numpy_hidden_dtype(hidden_cache_dtype) + autocast_enabled = bool(use_amp and device.type == "cuda") + + for batch in tqdm(loader, desc="Model/readout 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() + } + event_seq = batch_dev["event_seq"] + time_seq = batch_dev["time_seq"] + padding_mask = batch_dev["padding_mask"] + + amp_context = ( + torch.autocast(device_type=device.type, dtype=torch.float16) + if autocast_enabled else contextlib.nullcontext() + ) + with amp_context: + hidden = model( + event_seq=event_seq, + time_seq=time_seq, + sex=batch_dev["sex"], + padding_mask=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, + time_seq=time_seq, + padding_mask=padding_mask, + readout_mask=batch_dev["readout_mask"], + ) + + h = ro.hidden.detach().cpu().numpy().astype(out_dtype, copy=False) + hidden_parts.append(h) + max_len = max(max_len, h.shape[1]) + + for k in arrays: + if k == "sex": + arrays[k].append( + batch[k].cpu().numpy().astype(np.int8, copy=False)) + else: + arrays[k].append(batch[k].cpu().numpy()) + arrays["readout_mask"][-1] = ro.readout_mask.detach().cpu().numpy() + + def pad_3d(parts: List[np.ndarray], fill: float = 0.0) -> np.ndarray: + out = np.full( + (sum(x.shape[0] for x in parts), max_len, parts[0].shape[2]), + fill, + dtype=out_dtype, + ) + s = 0 + for x in parts: + out[s:s + x.shape[0], :x.shape[1], :] = x + s += x.shape[0] + return out + + def pad_2d(parts: List[np.ndarray], fill: Any = 0, dtype: Optional[np.dtype] = None) -> np.ndarray: + dtype = dtype or parts[0].dtype + out = np.full((sum(x.shape[0] + for x in parts), max_len), fill, dtype=dtype) + s = 0 + for x in parts: + out[s:s + x.shape[0], :x.shape[1]] = x + s += x.shape[0] + return out + + hidden_all = pad_3d(hidden_parts) + arr_out = { + "event_seq": pad_2d(arrays["event_seq"], PAD_IDX, np.int64), + "time_seq": pad_2d(arrays["time_seq"], 0.0, np.float32), + "target_event_seq": pad_2d(arrays["target_event_seq"], PAD_IDX, np.int64), + "target_time_seq": pad_2d(arrays["target_time_seq"], 0.0, np.float32), + "padding_mask": pad_2d(arrays["padding_mask"], False, bool), + "readout_mask": pad_2d(arrays["readout_mask"], False, bool), + "sex": np.concatenate(arrays["sex"], axis=0), + } + return hidden_all, arr_out + + +@torch.inference_mode() +def compute_logits_for_disease_chunk( + model: DeepHealth, + hidden_all: np.ndarray, + disease_ids: Sequence[int], + device: torch.device, + logit_batch_size: int, + use_amp: bool, +) -> np.ndarray: + """Project cached hidden states to only the requested disease columns.""" + n = int(hidden_all.shape[0]) + logit_batch_size = max(1, int(logit_batch_size)) + disease_ids = [int(x) for x in disease_ids] + + compute_dtype = torch.float16 if ( + device.type == "cuda" and use_amp) else torch.float32 + weight = model.risk_head.weight[disease_ids].detach().to( + device=device, dtype=compute_dtype) + + parts: List[np.ndarray] = [] + for start in tqdm(range(0, n, logit_batch_size), desc="Risk-head projection", leave=False, dynamic_ncols=True): + end = min(start + logit_batch_size, n) + h = torch.from_numpy(hidden_all[start:end]).to( + device=device, dtype=compute_dtype, non_blocking=True) + logits = torch.matmul(h, weight.t()) + parts.append(logits.float().cpu().numpy().astype( + np.float32, copy=False)) + del h, logits + return np.concatenate(parts, axis=0) + + +# --------------------------------------------------------------------------- +# CPU-parallel calibration AUC +# --------------------------------------------------------------------------- + +_WORKER: Dict[str, Any] = {} + + +def _build_flat_eval_table( + p_sex: np.ndarray, + time_seq: np.ndarray, + target_time_seq: np.ndarray, + target_event_seq: np.ndarray, + padding_mask: np.ndarray, + readout_mask: np.ndarray, + offset: float, + valid_target_min_id: int, + age_groups: np.ndarray, +) -> Dict[str, np.ndarray]: + """ + Build a compact event-level table for AUC evaluation. + + The previous version kept 2D patient x position arrays and each disease task + repeatedly scanned them. This table keeps only eligible target occurrences: + - target event is a real disease/event target; + - a valid readout token exists at least `offset` years before target time; + - the prediction token itself is valid and readout-valid; + - predicted age falls inside the requested age brackets. + + The disease-specific logic is unchanged; this only changes the physical layout + of the data to make CPU multiprocessing cheaper and more cache-friendly. + """ + if len(age_groups) < 2: + raise ValueError("age_groups must contain at least two values") + age_start = float(age_groups[0]) + age_step = float(age_groups[1] - age_groups[0]) + n_age = int(len(age_groups)) + + # Raw valid-target table is used only to decide which patients are cases for + # a disease. This intentionally happens before offset/age filtering, matching + # the supplied Delphi logic: a patient who ever has disease k is not a control + # for k, even if that occurrence later lacks a valid prediction point. + raw_patient_idx, raw_target_idx = np.where( + target_event_seq > int(valid_target_min_id)) + raw_target_event = target_event_seq[raw_patient_idx, raw_target_idx].astype( + np.int32, copy=False) + raw_sort_order = np.argsort(raw_target_event, kind="mergesort") + raw_sorted_target_event = raw_target_event[raw_sort_order] + raw_patient_idx = raw_patient_idx.astype(np.int32, copy=False) + + valid_pred_pos = ( + (time_seq[:, :, None] <= (target_time_seq[:, None, :] - float(offset))) + & readout_mask[:, :, None] + ) + pos_index = np.arange(time_seq.shape[1], dtype=np.int32)[None, :, None] + pred_idx_precompute = np.where( + valid_pred_pos, pos_index, -1).max(axis=1).astype(np.int32) + + candidate = (target_event_seq > int(valid_target_min_id)) & ( + pred_idx_precompute >= 0) + patient_idx, target_idx = np.where(candidate) + if patient_idx.size == 0: + return { + "patient": np.empty(0, dtype=np.int32), + "target_event": np.empty(0, dtype=np.int32), + "pred_idx": np.empty(0, dtype=np.int32), + "age_bin": np.empty(0, dtype=np.int16), + "target_time": np.empty(0, dtype=np.float32), + "sort_order": np.empty(0, dtype=np.int64), + "sorted_target_event": np.empty(0, dtype=np.int32), + "raw_patient": raw_patient_idx, + "raw_sort_order": raw_sort_order.astype(np.int64, copy=False), + "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False), + "p_sex": p_sex, + "age_groups": age_groups.astype(np.float32, copy=False), + "n_patients": np.int32(time_seq.shape[0]), + } + + pred_idx = pred_idx_precompute[patient_idx, target_idx] + pred_ok = padding_mask[patient_idx, + pred_idx] & readout_mask[patient_idx, pred_idx] + if not np.any(pred_ok): + return { + "patient": np.empty(0, dtype=np.int32), + "target_event": np.empty(0, dtype=np.int32), + "pred_idx": np.empty(0, dtype=np.int32), + "age_bin": np.empty(0, dtype=np.int16), + "target_time": np.empty(0, dtype=np.float32), + "sort_order": np.empty(0, dtype=np.int64), + "sorted_target_event": np.empty(0, dtype=np.int32), + "raw_patient": raw_patient_idx, + "raw_sort_order": raw_sort_order.astype(np.int64, copy=False), + "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False), + "p_sex": p_sex, + "age_groups": age_groups.astype(np.float32, copy=False), + "n_patients": np.int32(time_seq.shape[0]), + } + + patient_idx = patient_idx[pred_ok].astype(np.int32, copy=False) + target_idx = target_idx[pred_ok] + pred_idx = pred_idx[pred_ok].astype(np.int32, copy=False) + + pred_age = time_seq[patient_idx, pred_idx].astype(np.float32, copy=False) + age_bin = np.floor((pred_age - age_start) / age_step).astype(np.int16) + age_ok = (age_bin >= 0) & (age_bin < n_age) + if not np.any(age_ok): + return { + "patient": np.empty(0, dtype=np.int32), + "target_event": np.empty(0, dtype=np.int32), + "pred_idx": np.empty(0, dtype=np.int32), + "age_bin": np.empty(0, dtype=np.int16), + "target_time": np.empty(0, dtype=np.float32), + "sort_order": np.empty(0, dtype=np.int64), + "sorted_target_event": np.empty(0, dtype=np.int32), + "raw_patient": raw_patient_idx, + "raw_sort_order": raw_sort_order.astype(np.int64, copy=False), + "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False), + "p_sex": p_sex, + "age_groups": age_groups.astype(np.float32, copy=False), + "n_patients": np.int32(time_seq.shape[0]), + } + + patient_idx = patient_idx[age_ok] + target_idx = target_idx[age_ok] + pred_idx = pred_idx[age_ok] + age_bin = age_bin[age_ok] + target_event = target_event_seq[patient_idx, + target_idx].astype(np.int32, copy=False) + target_time = target_time_seq[patient_idx, + target_idx].astype(np.float32, copy=False) + + sort_order = np.argsort(target_event, kind="mergesort") + sorted_target_event = target_event[sort_order] + + return { + "patient": patient_idx.astype(np.int32, copy=False), + "target_event": target_event, + "pred_idx": pred_idx.astype(np.int32, copy=False), + "age_bin": age_bin.astype(np.int16, copy=False), + "target_time": target_time, + "sort_order": sort_order.astype(np.int64, copy=False), + "sorted_target_event": sorted_target_event.astype(np.int32, copy=False), + "raw_patient": raw_patient_idx, + "raw_sort_order": raw_sort_order.astype(np.int64, copy=False), + "raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False), + "p_sex": p_sex, + "age_groups": age_groups.astype(np.float32, copy=False), + "n_patients": np.int32(time_seq.shape[0]), + } + + +def _init_auc_worker_flat( + patient: np.ndarray, + target_event: np.ndarray, + pred_idx: np.ndarray, + age_bin: np.ndarray, + target_time: np.ndarray, + sort_order: np.ndarray, + sorted_target_event: np.ndarray, + raw_patient: np.ndarray, + raw_sort_order: np.ndarray, + raw_sorted_target_event: np.ndarray, + p_sex: np.ndarray, + age_groups: np.ndarray, + n_patients: int, + use_delong: bool, +): + # Prevent BLAS/OpenMP oversubscription when many worker processes are active. + 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") + + _WORKER.clear() + _WORKER.update({ + "patient": patient, + "target_event": target_event, + "pred_idx": pred_idx, + "age_bin": age_bin, + "target_time": target_time, + "sort_order": sort_order, + "sorted_target_event": sorted_target_event, + "raw_patient": raw_patient, + "raw_sort_order": raw_sort_order, + "raw_sorted_target_event": raw_sorted_target_event, + "p_sex": p_sex, + "age_groups": age_groups, + "n_patients": int(n_patients), + "use_delong": bool(use_delong), + }) + + +def _case_indices_for_token(token: int) -> np.ndarray: + sorted_target_event = _WORKER["sorted_target_event"] + sort_order = _WORKER["sort_order"] + left = np.searchsorted(sorted_target_event, int(token), side="left") + right = np.searchsorted(sorted_target_event, int(token), side="right") + if right <= left: + return np.empty(0, dtype=np.int64) + return sort_order[left:right] + + +def _raw_case_patients_for_token(token: int) -> np.ndarray: + raw_sorted_target_event = _WORKER["raw_sorted_target_event"] + raw_sort_order = _WORKER["raw_sort_order"] + raw_patient = _WORKER["raw_patient"] + left = np.searchsorted(raw_sorted_target_event, int(token), side="left") + right = np.searchsorted(raw_sorted_target_event, int(token), side="right") + if right <= left: + return np.empty(0, dtype=np.int32) + return np.unique(raw_patient[raw_sort_order[left:right]]) + + +def _calibration_auc_one_disease_flat(task: Tuple[int, int]) -> List[Dict[str, Any]]: + j, token = task + patient = _WORKER["patient"] + pred_idx = _WORKER["pred_idx"] + age_bin = _WORKER["age_bin"] + target_time = _WORKER["target_time"] + p_sex = _WORKER["p_sex"] + age_groups = _WORKER["age_groups"] + n_patients = _WORKER["n_patients"] + use_delong = _WORKER["use_delong"] + + case_idx = _case_indices_for_token(int(token)) + if case_idx.size < 2: + return [] + + case_patients = _raw_case_patients_for_token(int(token)) + if case_patients.size == 0: + return [] + + patient_has_case = np.zeros(n_patients, dtype=bool) + patient_has_case[case_patients] = True + + # Controls follow the supplied Delphi logic: any eligible target occurrence from + # a patient who never has this disease token in the evaluated target table. + control_idx = np.flatnonzero(~patient_has_case[patient]) + if control_idx.size == 0: + return [] + + out: List[Dict[str, Any]] = [] + for b, aa in enumerate(age_groups): + case_b = case_idx[age_bin[case_idx] == b] + control_b = control_idx[age_bin[control_idx] == b] + if case_b.size == 0 or control_b.size == 0: + continue + + # Match previous deterministic one-occurrence-per-patient behavior within + # each age bracket, separately for cases and controls. This avoids letting + # high-utilization patients dominate the AUC. + _, case_first = np.unique(patient[case_b], return_index=True) + _, control_first = np.unique(patient[control_b], return_index=True) + case_keep = case_b[case_first] + control_keep = control_b[control_first] + if case_keep.size == 0 or control_keep.size == 0: + continue + + # Delphi2M-aligned AUC score: use disease-specific eta/logit only. + # Prediction offset filters eligible prediction tokens but does not enter the score. + case_scores = p_sex[patient[case_keep], pred_idx[case_keep], j].astype( + np.float64, copy=False) + control_scores = p_sex[patient[control_keep], pred_idx[control_keep], j].astype( + np.float64, copy=False) + if case_scores.size == 0 or control_scores.size == 0: + continue + + if use_delong: + auc_value, auc_var = get_auc_delong_var( + control_scores, case_scores) + else: + auc_value, auc_var = get_auc_delong_var( + control_scores, case_scores) + + out.append({ + "token": int(token), + "auc": float(auc_value), + "auc_delong": float(auc_value), + "auc_variance_delong": float(auc_var), + "age": float(aa), + "age_right": float(aa + (age_groups[1] - age_groups[0])), + "n_healthy": int(control_scores.size), + "n_diseased": int(case_scores.size), + "mean_target_time": float(np.mean(target_time[case_keep])) if case_keep.size else np.nan, + }) + return out + + +def _calibration_auc_task_block(tasks: Sequence[Tuple[int, int]]) -> List[Dict[str, Any]]: + rows: List[Dict[str, Any]] = [] + for task in tasks: + rows.extend(_calibration_auc_one_disease_flat(task)) + return rows + + +def _split_tasks_for_workers( + tasks: Sequence[Tuple[int, int]], + effective_workers: int, + task_chunk_size: int, +) -> List[List[Tuple[int, int]]]: + if not tasks: + return [] + if task_chunk_size <= 0: + # Enough chunks to keep workers busy without creating one Future per disease. + task_chunk_size = max(1, math.ceil( + len(tasks) / max(1, effective_workers * 4))) + return [list(tasks[i:i + task_chunk_size]) for i in range(0, len(tasks), task_chunk_size)] + + +def compute_auc_chunk_parallel( + p_chunk: np.ndarray, + arrays: Dict[str, np.ndarray], + disease_ids: Sequence[int], + sex_value: int, + sex_name: str, + age_groups: np.ndarray, + offset: float, + valid_target_min_id: int, + num_workers: int, + use_delong: bool, + auc_task_chunk_size: int = 0, +) -> List[Dict[str, Any]]: + sex_mask = arrays["sex"] == sex_value + if not np.any(sex_mask): + return [] + + flat = _build_flat_eval_table( + p_sex=p_chunk[sex_mask], + time_seq=arrays["time_seq"][sex_mask], + target_time_seq=arrays["target_time_seq"][sex_mask], + target_event_seq=arrays["target_event_seq"][sex_mask], + padding_mask=arrays["padding_mask"][sex_mask], + readout_mask=arrays["readout_mask"][sex_mask], + offset=offset, + valid_target_min_id=valid_target_min_id, + age_groups=age_groups, + ) + if flat["patient"].size == 0: + return [] + + # Skip diseases with no cases in this sex before sending tasks to workers. + sorted_events = flat["sorted_target_event"] + tasks = [] + for j, token in enumerate(disease_ids): + left = np.searchsorted(sorted_events, int(token), side="left") + right = np.searchsorted(sorted_events, int(token), side="right") + if right - left >= 2: + tasks.append((j, int(token))) + if not tasks: + return [] + + effective_workers = max(1, min(int(num_workers), len(tasks))) + if effective_workers <= 1: + _init_auc_worker_flat( + flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"], + flat["target_time"], flat["sort_order"], flat["sorted_target_event"], + flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"], + flat["p_sex"], flat["age_groups"], int( + flat["n_patients"]), use_delong, + ) + nested = [_calibration_auc_one_disease_flat(t) for t in tqdm( + tasks, desc=f"AUC {sex_name}", leave=False, dynamic_ncols=True)] + else: + ctx = mp.get_context("fork") if hasattr( + os, "fork") else mp.get_context() + task_blocks = _split_tasks_for_workers( + tasks, effective_workers, int(auc_task_chunk_size)) + with ProcessPoolExecutor( + max_workers=effective_workers, + mp_context=ctx, + initializer=_init_auc_worker_flat, + initargs=( + flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"], + flat["target_time"], flat["sort_order"], flat["sorted_target_event"], + flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"], + flat["p_sex"], flat["age_groups"], int( + flat["n_patients"]), use_delong, + ), + ) as ex: + nested = list(tqdm( + ex.map(_calibration_auc_task_block, task_blocks), + total=len(task_blocks), + desc=f"AUC {sex_name}", + leave=False, + dynamic_ncols=True, + )) + + out: List[Dict[str, Any]] = [] + for rows in nested: + for r in rows: + r["sex"] = sex_name + r["offset"] = float(offset) + out.append(r) + return out + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- + +def evaluate_auc_pipeline( + model: DeepHealth, + loader: DataLoader, + dataset: HealthDataset, + output_path: Optional[str], + diseases_of_interest: Optional[Sequence[int]], + filter_min_total: int, + first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], + include_death: bool, + exclude_death: bool, + disease_chunk_size: int, + age_groups: np.ndarray, + offsets: Sequence[float], + device: torch.device, + attn_mask_mode: str, + readout_name: str, + readout_reduce: str, + num_workers_auc: int, + use_amp: bool, + auc_task_chunk_size: int = 0, + hidden_cache_dtype: str = "float16", + logit_batch_size: int = 256, +) -> Tuple[pd.DataFrame, pd.DataFrame]: + model.eval().to(device) + + disease_ids = select_disease_tokens( + dataset=dataset, + requested_tokens=diseases_of_interest, + filter_min_total=filter_min_total, + first_occurrence_by_token=first_occurrence_by_token, + ) + disease_ids = [int(k) for k in disease_ids if 0 <= + int(k) < dataset.vocab_size] + + death_token_ids = _get_death_token_ids(dataset) + if (not bool(include_death)) or bool(exclude_death): + before = len(disease_ids) + death_set = set(int(x) for x in death_token_ids) + disease_ids = [int(x) for x in disease_ids if int(x) not in death_set] + print( + "[INFO] Death exclusion applied on final disease_ids: " + f"include_death={bool(include_death)}, exclude_death={bool(exclude_death)}, " + f"before={before}, after={len(disease_ids)}.") + + if not disease_ids: + raise ValueError("No diseases selected for evaluation.") + + if disease_chunk_size is None or int(disease_chunk_size) <= 0: + disease_chunk_size = len(disease_ids) + disease_chunk_size = max(1, int(disease_chunk_size)) + + num_chunks = math.ceil(len(disease_ids) / disease_chunk_size) + chunks = np.array_split(np.asarray( + disease_ids, dtype=np.int64), num_chunks) + print( + f"Evaluating {len(disease_ids)} disease tokens in {len(chunks)} chunk(s).") + print("Using Delphi2M-aligned rate/logit score for AUC.") + print("AUC score = disease-specific eta at the latest eligible prediction token.") + print("Prediction offset controls eligibility only and does not enter the score.") + print(f"Evaluating prediction offsets: {', '.join(f'{x:g}' for x in offsets)} years.") + + # In current dataset sex is normalized to 0/1. UKB convention after normalization: 0=female, 1=male. + sex_items = [("female", 0), ("male", 1)] + all_rows: List[Dict[str, Any]] = [] + + valid_target_min_id = CHECKUP_IDX if NO_EVENT_IDX >= dataset.vocab_size else CHECKUP_IDX + # If NO_EVENT exists and should not be a disease/control target, require target > NO_EVENT_IDX. + if NO_EVENT_IDX in dataset.label_id_to_code and dataset.label_id_to_code.get(NO_EVENT_IDX) == "": + valid_target_min_id = NO_EVENT_IDX + + hidden_all, arrays = infer_readout_hidden( + model=model, + loader=loader, + device=device, + attn_mask_mode=attn_mask_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + use_amp=use_amp, + hidden_cache_dtype=hidden_cache_dtype, + ) + print( + f"Cached readout hidden: shape={hidden_all.shape}, dtype={hidden_all.dtype}") + + for chunk_idx, chunk in enumerate(tqdm(chunks, desc="Processing disease chunks", dynamic_ncols=True)): + p_chunk = compute_logits_for_disease_chunk( + model=model, + hidden_all=hidden_all, + disease_ids=chunk.tolist(), + device=device, + logit_batch_size=logit_batch_size, + use_amp=use_amp, + ) + for offset in offsets: + for sex_name, sex_value in sex_items: + rows = compute_auc_chunk_parallel( + p_chunk=p_chunk, + arrays=arrays, + disease_ids=chunk.tolist(), + sex_value=sex_value, + sex_name=sex_name, + age_groups=age_groups, + offset=float(offset), + valid_target_min_id=valid_target_min_id, + num_workers=num_workers_auc, + use_delong=True, + auc_task_chunk_size=auc_task_chunk_size, + ) + for r in rows: + r["disease_chunk_idx"] = int(chunk_idx) + all_rows.extend(rows) + del p_chunk + + del hidden_all, arrays + df_auc_unpooled = pd.DataFrame(all_rows) + if df_auc_unpooled.empty: + raise RuntimeError( + "No AUC rows were produced. Check offset, age_groups, eval split, and disease ids.") + + # Keep outputs self-contained with only evaluation fields and token->code mapping. + df_auc_unpooled["label_code"] = df_auc_unpooled["token"].map( + dataset.label_id_to_code) + + print("Using DeLong method to calculate AUC confidence intervals.") + grouped = df_auc_unpooled.groupby( + ["token", "label_code", "offset"], dropna=False, as_index=False) + df_auc = grouped.agg( + auc=("auc_delong", "mean"), + n_strata=("auc_delong", "size"), + n_diseased=("n_diseased", "sum"), + n_healthy=("n_healthy", "sum"), + auc_variance_sum=("auc_variance_delong", "sum"), + ) + df_auc["auc_variance_delong"] = ( + df_auc["auc_variance_sum"] + / (df_auc["n_strata"].clip(lower=1).astype(np.float64) ** 2) + ) + df_auc = df_auc.drop(columns=["auc_variance_sum"]) + + if output_path is not None: + out_dir = Path(output_path) + out_dir.mkdir(parents=True, exist_ok=True) + df_auc.to_csv(out_dir / "df_both.csv", index=False) + df_auc_unpooled.to_csv( + out_dir / "df_auc_unpooled.csv", index=False) + + return df_auc_unpooled, df_auc + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def parse_int_list(s: Any) -> Optional[List[int]]: + if s is None: + return None + if isinstance(s, (list, tuple, np.ndarray)): + return [int(x) for x in s] + text = str(s).strip() + if text == "": + return None + return [int(x.strip()) for x in text.split(",") if x.strip()] + + +def parse_float_list(s: Any) -> Optional[List[float]]: + if s is None: + return None + if isinstance(s, (list, tuple, np.ndarray)): + return [float(x) for x in s] + text = str(s).strip() + if text == "": + return None + return [float(x.strip()) for x in text.split(",") if x.strip()] + + +def make_auc_offsets(args: argparse.Namespace, cfg: Dict[str, Any]) -> List[float]: + explicit_offsets = parse_float_list(cfg_get(args, cfg, "offsets", None)) + if explicit_offsets is not None: + base_offsets = explicit_offsets + else: + next_token_offset = float(cfg_get(args, cfg, "offset", 0.1)) + base_offsets = [next_token_offset, 1.0, 5.0, 10.0] + + offsets: List[float] = [] + seen = set() + for value in base_offsets: + value = float(value) + key = round(value, 10) + if key not in seen: + offsets.append(value) + seen.add(key) + if not offsets: + raise ValueError("At least one AUC offset is required.") + return offsets + + +def main() -> None: + parser = argparse.ArgumentParser(description="Evaluate DeepHealth AUC") + + # Simplify arguments to only include run_path and output_path + parser.add_argument("--run_path", type=str, required=True, + help="Path containing train_config.json and best_model.pt") + parser.add_argument("--output_path", type=str, + default=None, help="Defaults to run_path") + parser.add_argument("--eval_split", type=str, default=None, + choices=["train", "val", "valid", + "validation", "test", "all"], + help="Evaluation split. Defaults to 'test' unless cfg contains eval_split.") + parser.add_argument("--dataset_subset_size", type=int, default=None, + help="Optional number of patients from the selected split.") + parser.add_argument("--batch_size", type=int, default=None, + help="Inference batch size; overrides train_config.json.") + parser.add_argument("--num_workers", type=int, default=None, + help="DataLoader workers; overrides train_config.json.") + parser.add_argument("--num_workers_auc", type=int, default=None, + help="CPU processes for AUC computation.") + parser.add_argument("--auc_task_chunk_size", type=int, default=None, + help="Diseases per submitted CPU task block. 0/None auto-tunes.") + parser.add_argument("--hidden_cache_dtype", type=str, default=None, choices=["float16", "float32"], + help="CPU dtype for cached readout hidden states. float16 saves memory and is usually enough for AUC.") + parser.add_argument("--logit_batch_size", type=int, default=None, + help="Patient batch size for projecting cached hidden states to disease logits.") + parser.add_argument("--disease_chunk_size", type=int, default=None, + help="Number of disease logits to materialize per inference pass. <=0 means one chunk (all diseases).") + parser.add_argument("--filter_min_total", type=int, default=None, + help="Minimum metadata count for disease selection; default 0.") + parser.add_argument("--offset", type=float, default=None, + help="Next-token prediction offset in years; preserved and evaluated alongside 1, 5, and 10 years by default.") + parser.add_argument("--offsets", type=str, default=None, + help="Comma-separated prediction offsets in years. Overrides the default set of offset,1,5,10.") + parser.add_argument("--age_start", type=float, default=None) + parser.add_argument("--age_stop", type=float, default=None) + parser.add_argument("--age_step", type=float, default=None) + parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None, + help="Use CUDA autocast during inference.") + + args = parser.parse_args() + + # Extract paths from run_path + run_path = Path(args.run_path) + config_path = run_path / "train_config.json" + model_ckpt_path = run_path / "best_model.pt" + + if not config_path.exists(): + raise FileNotFoundError(f"train_config.json not found in {run_path}") + if not model_ckpt_path.exists(): + raise FileNotFoundError(f"best_model.pt not found in {run_path}") + + cfg = load_json_config(str(config_path)) + + if args.output_path is None: + args.output_path = str(run_path) + + # Load configurations from train_config.json + data_prefix = cfg.get("data_prefix", "ukb") + labels_file = cfg.get("labels_file", "labels.csv") + no_event_interval_years = cfg.get("no_event_interval_years", 5.0) + include_no_event = cfg.get("include_no_event_in_uts_target", False) + + target_mode = cfg.get("target_mode", "uts") + dist_mode_cfg = cfg.get("dist_mode", "exponential") + attn_mask_mode = cfg.get( + "attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware") + readout_name = cfg.get( + "readout_name", "same_time_group_end" if target_mode == "uts" else "token") + readout_reduce = cfg.get("readout_reduce", "mean") + + device = torch.device(cfg.get("device", "cuda") + if torch.cuda.is_available() else "cpu") + if device.type == "cuda": + torch.backends.cudnn.benchmark = True + + print("Loading dataset...") + dataset = HealthDataset( + data_prefix=data_prefix, + labels_file=labels_file, + no_event_interval_years=no_event_interval_years, + include_no_event_in_uts_target=include_no_event, + extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), + ) + + subset, subset_indices = make_eval_subset(dataset, args, cfg) + print(f"Dataset: {len(dataset)} samples, vocab_size={dataset.vocab_size}") + + loader = DataLoader( + subset, + batch_size=int(cfg_get(args, cfg, "batch_size", 128)), + shuffle=False, + collate_fn=collate_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, + ) + + print("Building/loading model...") + state_dict = load_checkpoint_state_dict( + str(model_ckpt_path), map_location="cpu") + dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict) + cfg = dict(cfg) + cfg["dist_mode"] = dist_mode + print(f"Resolved dist_mode for evaluation: {dist_mode}") + + model = build_model_from_dataset(args, cfg, dataset).to(device) + load_model_state(model, str(model_ckpt_path), + device, state_dict=state_dict) + model.eval() + + age_groups = np.arange( + float(cfg_get(args, cfg, "age_start", 40.0)), + float(cfg_get(args, cfg, "age_stop", 80.0)), + float(cfg_get(args, cfg, "age_step", 5.0)), + dtype=np.float32, + ) + disease_spec = cfg_get(args, cfg, "diseases_of_interest", None) + if disease_spec is None: + disease_spec = cfg.get("disease_tokens", None) + diseases = parse_int_list(disease_spec) + first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps( + dataset, subset_indices) + include_death = bool(cfg_get(args, cfg, "include_death", True)) + exclude_death = bool(cfg_get(args, cfg, "exclude_death", False)) + auc_offsets = make_auc_offsets(args, cfg) + + evaluate_auc_pipeline( + model=model, + loader=loader, + dataset=dataset, + output_path=args.output_path, + diseases_of_interest=diseases, + filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)), + first_occurrence_by_token=first_occurrence_by_token, + include_death=include_death, + exclude_death=exclude_death, + disease_chunk_size=int(cfg_get(args, cfg, "disease_chunk_size", 0)), + age_groups=age_groups, + offsets=auc_offsets, + device=device, + attn_mask_mode=attn_mask_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max( + 1, (os.cpu_count() or 2) - 1))), + use_amp=bool(cfg_get(args, cfg, "use_amp", False)), + auc_task_chunk_size=int(cfg_get(args, cfg, "auc_task_chunk_size", 0)), + hidden_cache_dtype=str( + cfg_get(args, cfg, "hidden_cache_dtype", "float16")), + logit_batch_size=int( + cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))), + ) + + +if __name__ == "__main__": + main() diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py new file mode 100644 index 0000000..a98cfce --- /dev/null +++ b/evaluate_auc_v2.py @@ -0,0 +1,1399 @@ +"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth. + +This script is intentionally strict and supports only: + DeepHealth + exponential distribution + no_event imputation. + +Landmark querying is implemented by inserting a token at landmark age. +No t_query interface is used. +""" +from __future__ import annotations + +import argparse +import contextlib +import json +import math +import multiprocessing as mp +import os +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path +from typing import Any, Dict, 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 HealthDataset +from models import DeepHealth +from readouts import build_readout +from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX + +SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} +_TARGET_AWARE_MODES = {"target_aware", "delphi2m", "d2m"} + + +def load_json_config(path: Path) -> Dict[str, Any]: + if not path.exists(): + return {} + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +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 parse_int_list(value: Any) -> Optional[List[int]]: + if value is None: + return None + if isinstance(value, (list, tuple, np.ndarray)): + return [int(x) for x in value] + text = str(value).strip() + if text == "": + return None + return [int(x.strip()) for x in text.split(",") if x.strip()] + + +def parse_float_list(value: Any) -> Optional[List[float]]: + if value is None: + return None + if isinstance(value, (list, tuple, np.ndarray)): + return [float(x) for x in value] + text = str(value).strip() + if text == "": + return None + return [float(x.strip()) for x in text.split(",") if x.strip()] + + +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}") + rng = np.random.RandomState(int(seed)) + idx = rng.permutation(int(n)) + n_train = int(n * train_ratio) + n_val = int(n * val_ratio) + return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:] + + +def make_eval_indices(dataset: HealthDataset, 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 = split_map[eval_split] + + 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 np.asarray(indices, dtype=np.int64) + + +def load_checkpoint_state_dict(checkpoint_path: Path, map_location: str | torch.device = "cpu") -> Dict[str, Any]: + payload = torch.load(str(checkpoint_path), map_location=map_location) + if isinstance(payload, dict) and "model" in payload: + payload = payload["model"] + elif isinstance(payload, dict) and "state_dict" in payload: + payload = payload["state_dict"] + if not isinstance(payload, dict): + raise TypeError(f"Unsupported checkpoint payload type: {type(payload)}") + return payload + + +def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str: + mode = str(cfg_dist_mode).lower() + has_rho_head = any(str(k).startswith("rho_head.") + for k in state_dict.keys()) + if has_rho_head: + return "weibull" + return mode if mode in {"exponential", "weibull"} else "exponential" + + +def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth: + return DeepHealth( + vocab_size=dataset.vocab_size, + n_embd=int(cfg_get(args, cfg, "n_embd", 120)), + n_head=int(cfg_get(args, cfg, "n_head", 10)), + n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)), + n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)), + n_types=dataset.n_types, + n_cont_types=dataset.n_cont_types, + n_categories=dataset.n_categories, + cont_type_ids=dataset.cont_type_ids, + n_bins=int(cfg_get(args, cfg, "n_bins", 16)), + target_mode="next_token", + time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), + dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), + dropout=float(cfg_get(args, cfg, "dropout", 0.0)), + ) + + +def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None: + missing, unexpected = model.load_state_dict(state_dict, strict=False) + if missing or unexpected: + print( + f"[WARN] load_state_dict strict=False: missing={missing[:10]}, unexpected={unexpected[:10]}") + + +# --------------------------------------------------------------------------- +# DeLong AUC utilities +# --------------------------------------------------------------------------- + + +def compute_midrank(x: np.ndarray) -> np.ndarray: + x = np.asarray(x, dtype=np.float64) + order = np.argsort(x) + sorted_x = x[order] + n = len(x) + ranks = np.zeros(n, dtype=np.float64) + i = 0 + while i < n: + j = i + while j < n and sorted_x[j] == sorted_x[i]: + j += 1 + ranks[i:j] = 0.5 * (i + j - 1) + i = j + out = np.empty(n, dtype=np.float64) + out[order] = ranks + 1.0 + return out + + +def fast_delong(predictions_sorted_transposed: np.ndarray, label_1_count: int) -> Tuple[np.ndarray, np.ndarray]: + predictions_sorted_transposed = np.asarray( + predictions_sorted_transposed, dtype=np.float64) + m = int(label_1_count) + n = int(predictions_sorted_transposed.shape[1] - m) + if m <= 0 or n <= 0: + return np.array([np.nan], dtype=np.float64), np.array([[np.nan]], dtype=np.float64) + + positive_examples = predictions_sorted_transposed[:, :m] + negative_examples = predictions_sorted_transposed[:, m:] + k = int(predictions_sorted_transposed.shape[0]) + + tx = np.empty((k, m), dtype=np.float64) + ty = np.empty((k, n), dtype=np.float64) + tz = np.empty((k, m + n), dtype=np.float64) + for r in range(k): + tx[r] = compute_midrank(positive_examples[r]) + ty[r] = compute_midrank(negative_examples[r]) + tz[r] = compute_midrank(predictions_sorted_transposed[r]) + + aucs = tz[:, :m].sum(axis=1) / m / n - float(m + 1.0) / 2.0 / n + + v01 = (tz[:, :m] - tx) / n + v10 = 1.0 - (tz[:, m:] - ty) / m + + if k == 1: + sx = np.var(v01[0], ddof=1) if m > 1 else 0.0 + sy = np.var(v10[0], ddof=1) if n > 1 else 0.0 + delong_cov = np.array([[sx / m + sy / n]], dtype=np.float64) + else: + sx = np.cov(v01, rowvar=True) if m > 1 else np.zeros( + (k, k), dtype=np.float64) + sy = np.cov(v10, rowvar=True) if n > 1 else np.zeros( + (k, k), dtype=np.float64) + delong_cov = np.atleast_2d(sx) / m + np.atleast_2d(sy) / n + + return aucs, delong_cov + + +def get_auc_delong_var(case_scores: np.ndarray, control_scores: np.ndarray) -> Tuple[float, float]: + case_scores = np.asarray(case_scores, dtype=np.float64) + control_scores = np.asarray(control_scores, dtype=np.float64) + if case_scores.size == 0 or control_scores.size == 0: + return np.nan, np.nan + + y_true = np.array([1] * len(case_scores) + [0] * + len(control_scores), dtype=np.int8) + y_score = np.concatenate([case_scores, control_scores]).astype( + np.float64, copy=False) + order = (-y_true).argsort() + + aucs, cov = fast_delong(y_score[np.newaxis, order], int(y_true.sum())) + var = float(np.asarray(cov).reshape(-1)[0]) + return float(aucs[0]), var + + +# --------------------------------------------------------------------------- +# Metadata and token selection +# --------------------------------------------------------------------------- + + +def _first_existing_column(df: pd.DataFrame, candidates: Sequence[str]) -> Optional[str]: + for col in candidates: + if col in df.columns: + return col + return None + + +def build_metadata_for_merge(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> pd.DataFrame: + base_rows = [] + for token, code in dataset.label_id_to_code.items(): + token = int(token) + code_text = str(code) + if token in SPECIAL_TOKENS or code_text.startswith("<"): + continue + base_rows.append({"token": token, "label_code": code_text}) + base = pd.DataFrame(base_rows) + if labels_meta is None or labels_meta.empty: + return base + + meta = labels_meta.copy() + code_col = _first_existing_column( + meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"]) + if code_col is not None: + meta["_label_code"] = meta[code_col].astype( + str).map(lambda s: s.split()[0].strip()) + merged = base.merge(meta, left_on="label_code", + right_on="_label_code", how="left") + return merged.drop(columns=["_label_code"], errors="ignore") + + if "index" in meta.columns: + idx = pd.to_numeric(meta["index"], errors="coerce") + has_no_event = ( + NO_EVENT_IDX in dataset.label_id_to_code + and dataset.label_id_to_code.get(NO_EVENT_IDX) == "" + ) + if has_no_event: + idx = idx.where(idx < NO_EVENT_IDX, idx + 1) + meta["_index_int"] = idx.astype("Int64") + merged = base.merge(meta, left_on="token", + right_on="_index_int", how="left") + return merged.drop(columns=["_index_int"], errors="ignore") + + return base + + +def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> Dict[int, float]: + if labels_meta is None or labels_meta.empty or "count" not in labels_meta.columns: + return {} + + out: Dict[int, float] = {} + meta = labels_meta.copy() + count_series = pd.to_numeric(meta["count"], errors="coerce") + + code_col = _first_existing_column( + meta, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"]) + if code_col is not None: + for code_text, count in zip(meta[code_col].astype(str).tolist(), count_series.tolist()): + code = code_text.split()[0].strip() + if code in dataset.label_code_to_id and pd.notna(count): + out[int(dataset.label_code_to_id[code])] = float(count) + if out: + return out + + if "index" in meta.columns: + idx = pd.to_numeric(meta["index"], errors="coerce") + has_no_event = ( + NO_EVENT_IDX in dataset.label_id_to_code + and dataset.label_id_to_code.get(NO_EVENT_IDX) == "" + ) + if has_no_event: + idx = idx.where(idx < NO_EVENT_IDX, idx + 1) + for token, count in zip(idx.tolist(), count_series.tolist()): + if pd.notna(token) and pd.notna(count): + out[int(token)] = float(count) + return out + + +def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]: + ids: List[int] = [] + + if labels_meta is not None and not labels_meta.empty: + meta = labels_meta.copy() + if "ICD-10 Chapter (short)" in meta.columns: + death_rows = meta[meta["ICD-10 Chapter (short)"].astype( + str) == "Death"] + code_col = _first_existing_column( + death_rows, ["Name", "code", "ICD10", "icd10", "label", "token", "disease_code"]) + if code_col is not None: + for raw in death_rows[code_col].astype(str).tolist(): + code = raw.split()[0].strip() + if code in dataset.label_code_to_id: + ids.append(int(dataset.label_code_to_id[code])) + elif "index" in death_rows.columns: + idx = pd.to_numeric(death_rows["index"], errors="coerce") + has_no_event = ( + NO_EVENT_IDX in dataset.label_id_to_code + and dataset.label_id_to_code.get(NO_EVENT_IDX) == "" + ) + if has_no_event: + idx = idx.where(idx < NO_EVENT_IDX, idx + 1) + ids.extend(int(x) for x in idx.dropna().astype(int).tolist()) + + exact_codes = {"death", "", "dth", "deceased", "mortality"} + for token, code in dataset.label_id_to_code.items(): + token = int(token) + if token in SPECIAL_TOKENS: + continue + text = str(code).strip().lower() + if text in exact_codes or ("death" in text) or ("mortality" in text): + ids.append(token) + + return sorted(set(int(x) for x in ids if int(x) not in SPECIAL_TOKENS)) + + +def _build_first_occurrence_maps( + dataset: HealthDataset, + subset_indices: np.ndarray, +) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]: + patient_count = len(subset_indices) + followup_end = np.full(patient_count, -np.inf, dtype=np.float32) + death_time = np.full(patient_count, np.inf, dtype=np.float32) + sex = np.full(patient_count, -1, dtype=np.int8) + + first_lists: Dict[int, List[Tuple[int, float]]] = {} + + for patient_id, dataset_index in enumerate(subset_indices.tolist()): + s = dataset.samples[int(dataset_index)] + + seq_event = np.asarray(s["event_seq"], dtype=np.int64) + seq_time = np.asarray(s["time_seq"], dtype=np.float32) + tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64) + tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32) + + if seq_event.size == 0 or tgt_event.size == 0: + continue + + full_event = np.concatenate([seq_event, tgt_event[-1:]]) + full_time = np.concatenate([seq_time, tgt_time[-1:]]) + + sex[patient_id] = int(s["sex"]) + followup_end[patient_id] = np.max(full_time).astype(np.float32) + + uniq_tokens, first_idx = np.unique(full_event, return_index=True) + for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()): + token = int(token) + event_time = float(full_time[int(idx)]) + if token not in first_lists: + first_lists[token] = [] + first_lists[token].append((patient_id, event_time)) + + packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {} + for token, pairs in first_lists.items(): + if not pairs: + continue + packed[int(token)] = ( + np.asarray([p for p, _ in pairs], dtype=np.int32), + np.asarray([t for _, t in pairs], dtype=np.float32), + ) + + return packed, followup_end, death_time, sex + + +def select_disease_tokens( + dataset: HealthDataset, + labels_meta: Optional[pd.DataFrame], + requested_tokens: Optional[Sequence[int]], + filter_min_total: int, + first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], +) -> List[int]: + base = [ + int(token) + for token, code in dataset.label_id_to_code.items() + if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<") + ] + base_set = set(base) + + if requested_tokens is not None: + return sorted(set(int(x) for x in requested_tokens if int(x) in base_set)) + + disease_ids = sorted(base) + if int(filter_min_total) <= 0: + return disease_ids + + meta_counts = _metadata_count_map(dataset, labels_meta) + if meta_counts: + return [token for token in disease_ids if float(meta_counts.get(token, 0.0)) > float(filter_min_total)] + + split_counts = {} + for token, pairs in first_occurrence_by_token.items(): + token = int(token) + if token not in base_set: + continue + split_counts[token] = len(np.unique(pairs[0])) + return [token for token in disease_ids if int(split_counts.get(token, 0)) > int(filter_min_total)] + + +# --------------------------------------------------------------------------- +# Landmark dataset +# --------------------------------------------------------------------------- + + +class LandmarkDataset(Dataset): + def __init__( + self, + dataset: HealthDataset, + subset_indices: np.ndarray, + landmark_ages: np.ndarray, + attn_mask_mode: str, + min_history_events: int, + first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], + death_token_ids: Sequence[int], + ) -> None: + self.dataset = dataset + self.subset_indices = np.asarray(subset_indices, dtype=np.int64) + self.landmark_ages = np.asarray(landmark_ages, dtype=np.float32) + self.attn_mask_mode = str(attn_mask_mode).lower() + self.min_history_events = int(min_history_events) + + self.first_occurrence_by_token = first_occurrence_by_token + self.death_token_ids = sorted(set(int(x) for x in death_token_ids)) + + rows: List[Dict[str, Any]] = [] + + self.patient_followup_end = np.full( + len(self.subset_indices), -np.inf, dtype=np.float32) + self.patient_death_time = np.full( + len(self.subset_indices), np.inf, dtype=np.float32) + self.patient_sex = np.full(len(self.subset_indices), -1, dtype=np.int8) + + for patient_id, dataset_index in enumerate(self.subset_indices.tolist()): + s = dataset.samples[int(dataset_index)] + + seq_event = np.asarray(s["event_seq"], dtype=np.int64) + seq_time = np.asarray(s["time_seq"], dtype=np.float32) + tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64) + tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32) + + if seq_event.size == 0 or tgt_event.size == 0: + continue + + full_event = np.concatenate([seq_event, tgt_event[-1:]]) + full_time = np.concatenate([seq_time, tgt_time[-1:]]) + + self.patient_sex[patient_id] = int(s["sex"]) + followup_end = float(np.max(full_time)) + self.patient_followup_end[patient_id] = np.float32(followup_end) + + d_time = np.inf + for death_token in self.death_token_ids: + hit = np.where(full_event == int(death_token))[0] + if hit.size > 0: + d_time = min(d_time, float(full_time[int(hit[0])])) + self.patient_death_time[patient_id] = np.float32(d_time) + + for landmark_age in self.landmark_ages.tolist(): + landmark_age = float(landmark_age) + + if not (followup_end > landmark_age): + continue + if not (float(self.patient_death_time[patient_id]) > landmark_age): + continue + + prefix_mask = full_time <= landmark_age + if not np.any(prefix_mask): + continue + + prefix_events = full_event[prefix_mask] + prefix_times = full_time[prefix_mask] + + valid_history_mask = ~np.isin(prefix_events, np.array( + [PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64)) + if valid_history_mask.sum() < self.min_history_events: + continue + + event_seq_landmark = np.concatenate( + [ + prefix_events.astype(np.int64, copy=False), + np.array([NO_EVENT_IDX], dtype=np.int64), + ] + ) + time_seq_landmark = np.concatenate( + [ + prefix_times.astype(np.float32, copy=False), + np.array([np.float32(landmark_age)], dtype=np.float32), + ] + ) + + target_time_seq = time_seq_landmark.copy() + if self.attn_mask_mode in _TARGET_AWARE_MODES: + target_time_seq[-1] = np.nextafter( + np.float32(landmark_age), np.float32(np.inf), dtype=np.float32 + ) + + readout_mask = np.zeros(len(event_seq_landmark), dtype=bool) + readout_mask[-1] = True + + rows.append( + { + "patient_id": int(patient_id), + "dataset_index": int(dataset_index), + "sex": int(s["sex"]), + "landmark_age": np.float32(landmark_age), + "followup_end_time": np.float32(followup_end), + "death_time": np.float32(self.patient_death_time[patient_id]), + "landmark_pos": int(len(event_seq_landmark) - 1), + "event_seq": event_seq_landmark, + "time_seq": time_seq_landmark, + "target_time_seq": target_time_seq, + "readout_mask": readout_mask, + "other_type": np.asarray(s["other_type"], dtype=np.int64), + "other_value": np.asarray(s["other_value"], dtype=np.float32), + "other_value_kind": np.asarray(s["other_value_kind"], dtype=np.int64), + "other_time": np.asarray(s["other_time"], dtype=np.float32), + } + ) + + if not rows: + raise RuntimeError( + "No eligible landmark query samples were produced. Check eval split, landmark ages, and min_history_events." + ) + + self.rows = rows + + def __len__(self) -> int: + return len(self.rows) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + s = self.rows[idx] + return { + "event_seq": torch.from_numpy(s["event_seq"]).long(), + "time_seq": torch.from_numpy(s["time_seq"]).float(), + "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(), + "readout_mask": torch.from_numpy(s["readout_mask"]), + "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(), + "landmark_pos": torch.tensor(s["landmark_pos"], dtype=torch.long), + "patient_id": torch.tensor(s["patient_id"], dtype=torch.long), + "landmark_age": torch.tensor(float(s["landmark_age"]), dtype=torch.float32), + "followup_end_time": torch.tensor(float(s["followup_end_time"]), dtype=torch.float32), + "death_time": torch.tensor(float(s["death_time"]), dtype=torch.float32), + } + + +def collate_landmark_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) + target_time_seq = pad_sequence( + [x["target_time_seq"] for x in batch], batch_first=True, padding_value=0.0) + readout_mask = pad_sequence( + [x["readout_mask"] for x in batch], batch_first=True, padding_value=False) + 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) + + return { + "event_seq": event_seq, + "time_seq": time_seq, + "target_time_seq": target_time_seq, + "padding_mask": event_seq > PAD_IDX, + "readout_mask": readout_mask, + "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, + "landmark_pos": torch.stack([x["landmark_pos"] for x in batch]), + "patient_id": torch.stack([x["patient_id"] for x in batch]), + "landmark_age": torch.stack([x["landmark_age"] for x in batch]), + "followup_end_time": torch.stack([x["followup_end_time"] for x in batch]), + "death_time": torch.stack([x["death_time"] for x in batch]), + } + + +def _numpy_hidden_dtype(name: str) -> np.dtype: + key = str(name).lower() + if key in {"float16", "fp16", "half", "bfloat16", "bf16"}: + return np.float16 + if key in {"float32", "fp32", "single"}: + return np.float32 + raise ValueError( + f"hidden_cache_dtype must be float16 or float32, got {name!r}") + + +@torch.inference_mode() +def infer_landmark_hidden( + model: DeepHealth, + loader: DataLoader, + device: torch.device, + attn_mask_mode: str, + readout_name: str, + readout_reduce: str, + use_amp: bool, + hidden_cache_dtype: str, +) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: + 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] = [] + arrays = { + "patient_id": [], + "sex": [], + "landmark_age": [], + "followup_end_time": [], + "death_time": [], + } + + out_dtype = _numpy_hidden_dtype(hidden_cache_dtype) + amp_enabled = bool(use_amp and device.type == "cuda") + + for batch in tqdm(loader, desc="Landmark 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_ctx = ( + torch.autocast(device_type=device.type, dtype=torch.float16) + if amp_enabled + else contextlib.nullcontext() + ) + with amp_ctx: + hidden = 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", + ) + readout_out = readout( + hidden=hidden, + time_seq=batch_dev["time_seq"], + padding_mask=batch_dev["padding_mask"], + readout_mask=batch_dev["readout_mask"], + ) + + landmark_hidden = readout_out.hidden.gather( + 1, + batch_dev["landmark_pos"].long()[:, None, None].expand(-1, + 1, readout_out.hidden.shape[-1]), + ).squeeze(1) + + hidden_parts.append(landmark_hidden.detach( + ).cpu().numpy().astype(out_dtype, copy=False)) + arrays["patient_id"].append( + batch["patient_id"].cpu().numpy().astype(np.int32, copy=False)) + arrays["sex"].append( + batch["sex"].cpu().numpy().astype(np.int8, copy=False)) + arrays["landmark_age"].append( + batch["landmark_age"].cpu().numpy().astype(np.float32, copy=False)) + arrays["followup_end_time"].append( + batch["followup_end_time"].cpu().numpy().astype(np.float32, copy=False)) + arrays["death_time"].append( + batch["death_time"].cpu().numpy().astype(np.float32, copy=False)) + + hidden_all = np.concatenate(hidden_parts, axis=0) + row_arrays = { + k: np.concatenate(v, axis=0) for k, v in arrays.items() + } + return hidden_all, row_arrays + + +@torch.inference_mode() +def project_logits_chunk( + model: DeepHealth, + hidden_all: np.ndarray, + disease_ids: Sequence[int], + device: torch.device, + logit_batch_size: int, + use_amp: bool, +) -> np.ndarray: + n = int(hidden_all.shape[0]) + logit_batch_size = max(1, int(logit_batch_size)) + disease_ids = [int(x) for x in disease_ids] + + compute_dtype = torch.float16 if ( + device.type == "cuda" and use_amp) else torch.float32 + weight = model.risk_head.weight[disease_ids].detach().to( + device=device, dtype=compute_dtype) + + out_parts: List[np.ndarray] = [] + for start in tqdm(range(0, n, logit_batch_size), desc="Risk projection", leave=False, dynamic_ncols=True): + end = min(start + logit_batch_size, n) + h = torch.from_numpy(hidden_all[start:end]).to( + device=device, dtype=compute_dtype, non_blocking=True) + logits = torch.matmul(h, weight.t()) + out_parts.append(logits.float().cpu( + ).numpy().astype(np.float32, copy=False)) + del h, logits + return np.concatenate(out_parts, axis=0) + + +# --------------------------------------------------------------------------- +# Parallel AUC workers +# --------------------------------------------------------------------------- + + +_WORKER: Dict[str, Any] = {} + + +def _init_worker( + disease_ids: np.ndarray, + score_chunk: np.ndarray, + row_patient_id: np.ndarray, + row_sex: np.ndarray, + row_landmark_age: np.ndarray, + row_followup_end: np.ndarray, + row_death_time: np.ndarray, + first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], + patient_count: int, + horizons: np.ndarray, + min_cases: int, + exclude_death_competing: bool, + death_token_ids: np.ndarray, +) -> 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") + + _WORKER.clear() + _WORKER.update( + { + "disease_ids": np.asarray(disease_ids, dtype=np.int64), + "score_chunk": np.asarray(score_chunk, dtype=np.float32), + "row_patient_id": np.asarray(row_patient_id, dtype=np.int32), + "row_sex": np.asarray(row_sex, dtype=np.int8), + "row_landmark_age": np.asarray(row_landmark_age, dtype=np.float32), + "row_followup_end": np.asarray(row_followup_end, dtype=np.float32), + "row_death_time": np.asarray(row_death_time, 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), + "exclude_death_competing": bool(exclude_death_competing), + "death_token_ids": set(int(x) for x in np.asarray(death_token_ids, dtype=np.int64).tolist()), + "first_time_cache": {}, + } + ) + + +def _first_time_by_patient(token: int) -> np.ndarray: + cache = _WORKER["first_time_cache"] + if int(token) in cache: + return cache[int(token)] + + arr = np.full(int(_WORKER["patient_count"]), np.inf, dtype=np.float32) + pairs = _WORKER["first_occurrence_by_token"].get(int(token)) + if pairs is not None: + p, t = pairs + arr[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32) + cache[int(token)] = arr + return arr + + +def _score_to_probability(logits: np.ndarray, score_mode: str, horizon: float) -> np.ndarray: + if score_mode == "eta": + return logits.astype(np.float64, copy=False) + rate = np.log1p(np.exp(-np.abs(logits))) + np.maximum(logits, 0.0) + rate = rate + np.float32(1e-8) + return (-np.expm1(-rate * np.float32(horizon))).astype(np.float64, copy=False) + + +def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]: + j, token, score_mode = task + + token = int(token) + row_patient_id = _WORKER["row_patient_id"] + row_sex = _WORKER["row_sex"] + row_landmark_age = _WORKER["row_landmark_age"] + row_followup_end = _WORKER["row_followup_end"] + row_death_time = _WORKER["row_death_time"] + logits_token = _WORKER["score_chunk"][:, int(j)] + + first_time_patient = _first_time_by_patient(token) + is_death_target = token in _WORKER["death_token_ids"] + horizons = _WORKER["horizons"] + + out_rows: List[Dict[str, Any]] = [] + + for sex_value, sex_name in [(0, "female"), (1, "male")]: + sex_mask_rows = row_sex == sex_value + if not np.any(sex_mask_rows): + continue + + lm_values = np.unique(row_landmark_age[sex_mask_rows]) + for a in lm_values.tolist(): + a = float(a) + stratum_mask = sex_mask_rows & np.isclose( + row_landmark_age, np.float32(a)) + if not np.any(stratum_mask): + continue + + idx = np.flatnonzero(stratum_mask) + pid = row_patient_id[idx] + followup = row_followup_end[idx] + death_time = row_death_time[idx] + first_time = first_time_patient[pid] + + prevalent = first_time <= np.float32(a) + if np.all(prevalent): + continue + + for horizon in horizons.tolist(): + horizon = float(horizon) + h_end = np.float32(a + horizon) + + cases = (first_time > np.float32(a)) & (first_time <= h_end) + controls = (~prevalent) & ( + ((~np.isinf(first_time)) & (first_time > h_end)) + | (np.isinf(first_time) & (followup >= h_end)) + ) + + if bool(_WORKER["exclude_death_competing"]) and (not is_death_target): + death_in_window = (death_time > np.float32(a)) & ( + death_time <= h_end) + death_before_disease = death_time < first_time + competing_death = death_in_window & death_before_disease + controls = controls & (~competing_death) + + case_idx = np.flatnonzero(cases) + control_idx = np.flatnonzero(controls) + if case_idx.size < int(_WORKER["min_cases"]) or control_idx.size == 0: + continue + + case_scores = _score_to_probability( + logits_token[idx[case_idx]], score_mode=score_mode, horizon=horizon) + control_scores = _score_to_probability( + logits_token[idx[control_idx]], score_mode=score_mode, horizon=horizon) + + auc, auc_var = get_auc_delong_var(case_scores, control_scores) + if np.isnan(auc) or np.isnan(auc_var): + continue + + out_rows.append( + { + "token": token, + "sex": sex_name, + "landmark_age": float(a), + "horizon": float(horizon), + "auc": float(auc), + "auc_delong": float(auc), + "auc_variance_delong": float(auc_var), + "n_diseased": int(case_scores.size), + "n_healthy": int(control_scores.size), + } + ) + + return out_rows + + +def _token_task_block(tasks: Sequence[Tuple[int, int, str]]) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + for t in tasks: + out.extend(_eval_token(t)) + return out + + +def _split_tasks(tasks: Sequence[Tuple[int, int, str]], workers: int, chunk_size: int) -> List[List[Tuple[int, int, str]]]: + if not tasks: + return [] + if int(chunk_size) <= 0: + chunk_size = max(1, math.ceil(len(tasks) / max(1, workers * 4))) + return [list(tasks[i:i + chunk_size]) for i in range(0, len(tasks), chunk_size)] + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- + + +def evaluate_landmark_auc( + model: DeepHealth, + loader: DataLoader, + landmark_dataset: LandmarkDataset, + output_path: Path, + labels_meta: Optional[pd.DataFrame], + disease_ids: Sequence[int], + disease_chunk_size: int, + score_mode: str, + horizons: np.ndarray, + device: torch.device, + attn_mask_mode: str, + readout_name: str, + readout_reduce: str, + num_workers_auc: int, + auc_task_chunk_size: int, + min_cases: int, + exclude_death_competing: bool, + use_amp: bool, + hidden_cache_dtype: str, + logit_batch_size: int, + meta_info: Dict[str, Any], +) -> Tuple[pd.DataFrame, pd.DataFrame]: + model.eval().to(device) + + hidden_all, row_arrays = infer_landmark_hidden( + model=model, + loader=loader, + device=device, + attn_mask_mode=attn_mask_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + use_amp=use_amp, + hidden_cache_dtype=hidden_cache_dtype, + ) + print( + f"Cached landmark hidden: shape={hidden_all.shape}, dtype={hidden_all.dtype}") + + disease_ids = [int(x) for x in disease_ids] + if disease_chunk_size <= 0: + disease_chunk_size = len(disease_ids) + chunks = np.array_split(np.asarray(disease_ids, dtype=np.int64), math.ceil( + len(disease_ids) / disease_chunk_size)) + + all_rows: List[Dict[str, Any]] = [] + for chunk_idx, chunk in enumerate(tqdm(chunks, desc="Disease chunks", dynamic_ncols=True)): + chunk_ids = chunk.tolist() + logits_chunk = project_logits_chunk( + model=model, + hidden_all=hidden_all, + disease_ids=chunk_ids, + device=device, + logit_batch_size=logit_batch_size, + use_amp=use_amp, + ) + + tasks = [(j, int(token), score_mode) + for j, token in enumerate(chunk_ids)] + workers = max(1, min(int(num_workers_auc), len(tasks))) + + if workers <= 1: + _init_worker( + disease_ids=np.asarray(chunk_ids, dtype=np.int64), + score_chunk=logits_chunk, + row_patient_id=row_arrays["patient_id"], + row_sex=row_arrays["sex"], + row_landmark_age=row_arrays["landmark_age"], + row_followup_end=row_arrays["followup_end_time"], + row_death_time=row_arrays["death_time"], + first_occurrence_by_token=landmark_dataset.first_occurrence_by_token, + patient_count=len(landmark_dataset.subset_indices), + horizons=horizons, + min_cases=min_cases, + exclude_death_competing=exclude_death_competing, + death_token_ids=np.asarray( + landmark_dataset.death_token_ids, dtype=np.int64), + ) + nested = [_eval_token(t) for t in tqdm( + tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)] + else: + ctx = mp.get_context("spawn") + blocks = _split_tasks(tasks, workers, auc_task_chunk_size) + with ProcessPoolExecutor( + max_workers=workers, + mp_context=ctx, + initializer=_init_worker, + initargs=( + np.asarray(chunk_ids, dtype=np.int64), + logits_chunk, + row_arrays["patient_id"], + row_arrays["sex"], + row_arrays["landmark_age"], + row_arrays["followup_end_time"], + row_arrays["death_time"], + landmark_dataset.first_occurrence_by_token, + len(landmark_dataset.subset_indices), + horizons, + min_cases, + exclude_death_competing, + np.asarray(landmark_dataset.death_token_ids, + dtype=np.int64), + ), + ) as ex: + nested = list( + tqdm( + ex.map(_token_task_block, blocks), + total=len(blocks), + desc=f"AUC chunk {chunk_idx}", + leave=False, + dynamic_ncols=True, + ) + ) + + for rows in nested: + for r in rows: + r["disease_chunk_idx"] = int(chunk_idx) + all_rows.append(r) + + del logits_chunk + + if not all_rows: + raise RuntimeError( + "No AUC rows produced. Check landmark ages, horizons, min_cases, follow-up, or disease token selection." + ) + + df_unpooled = pd.DataFrame(all_rows) + df_unpooled["label_code"] = df_unpooled["token"].map( + landmark_dataset.dataset.label_id_to_code) + + for k, v in meta_info.items(): + df_unpooled[k] = v + + meta_table = build_metadata_for_merge(landmark_dataset.dataset, labels_meta) + df_unpooled = df_unpooled.merge( + meta_table, on=["token", "label_code"], how="left") + + grouped = df_unpooled.groupby( + ["token", "label_code", "horizon"], dropna=False, as_index=False) + df_merged = grouped.agg( + auc=("auc_delong", "mean"), + n_strata=("auc_delong", "size"), + n_diseased=("n_diseased", "sum"), + n_healthy=("n_healthy", "sum"), + auc_variance_sum=("auc_variance_delong", "sum"), + ) + df_merged["auc_variance_delong"] = ( + df_merged["auc_variance_sum"] + / (df_merged["n_strata"].clip(lower=1).astype(np.float64) ** 2) + ) + df_merged = df_merged.drop(columns=["auc_variance_sum"]) + + keep_meta = [ + c for c in [ + "model_ckpt_path", + "config_path", + "target_mode", + "dist_mode", + "time_mode", + "attn_mask_mode", + "readout_name", + "landmark_query_mode", + "landmark_token_mode", + "score_mode", + "eval_split", + ] + if c in df_unpooled.columns + ] + for col in keep_meta: + df_merged[col] = meta_info[col] + + output_path.mkdir(parents=True, exist_ok=True) + df_unpooled.to_csv( + output_path / "df_auc_landmark_unpooled.csv", index=False) + df_merged.to_csv(output_path / "df_auc_landmark.csv", index=False) + + return df_unpooled, df_merged + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Evaluate DeepHealth landmark fixed-horizon incident 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="test", + 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("--disease_chunk_size", type=int, default=None) + parser.add_argument("--filter_min_total", type=int, default=None) + parser.add_argument("--labels_meta_path", type=str, default=None) + + parser.add_argument("--landmark_start", type=float, default=None) + parser.add_argument("--landmark_stop", type=float, default=None) + parser.add_argument("--landmark_step", type=float, default=None) + parser.add_argument("--horizons", type=str, default=None) + + parser.add_argument("--min_cases", type=int, default=None) + parser.add_argument("--min_history_events", type=int, default=None) + + parser.add_argument("--score_mode", type=str, + default=None, choices=["risk", "eta"]) + parser.add_argument("--exclude_death_in_window_without_disease", + action=argparse.BooleanOptionalAction, default=None) + parser.add_argument( + "--use_amp", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--logit_batch_size", type=int, default=None) + parser.add_argument("--hidden_cache_dtype", type=str, + default=None, choices=["float16", "float32"]) + + args = parser.parse_args() + + run_path = Path(args.run_path) + config_path = run_path / "train_config.json" + model_ckpt_path = run_path / "best_model.pt" + + if not config_path.exists(): + raise FileNotFoundError(f"train_config.json not found in {run_path}") + if not model_ckpt_path.exists(): + raise FileNotFoundError(f"best_model.pt not found in {run_path}") + + cfg = load_json_config(config_path) + + data_prefix = cfg.get("data_prefix", "ukb") + labels_file = cfg.get("labels_file", "labels.csv") + no_event_interval_years = cfg.get("no_event_interval_years", 5.0) + include_no_event_in_uts_target = cfg.get( + "include_no_event_in_uts_target", False) + + target_mode = cfg.get("target_mode", "uts") + dist_mode_cfg = str(cfg.get("dist_mode", "exponential")) + attn_mask_mode = str(cfg.get( + "attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")) + 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")) + time_mode = str(cfg.get("time_mode", "relative")) + + output_path = Path( + cfg_get(args, cfg, "output_path", None) + or cfg.get("output_path", None) + or str(run_path) + ) + output_path.mkdir(parents=True, exist_ok=True) + + 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 = None + if labels_meta_path is not None and Path(str(labels_meta_path)).exists(): + labels_meta = pd.read_csv(str(labels_meta_path)) + + print("Loading dataset...") + dataset = HealthDataset( + data_prefix=data_prefix, + labels_file=labels_file, + no_event_interval_years=float(no_event_interval_years), + include_no_event_in_uts_target=bool(include_no_event_in_uts_target), + extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), + ) + + has_no_event = ( + NO_EVENT_IDX in dataset.label_id_to_code + and dataset.label_id_to_code.get(NO_EVENT_IDX) == "" + and dataset.vocab_size > NO_EVENT_IDX + ) + if not has_no_event: + print( + "[SKIP] This checkpoint/run does not support imputation. " + "Landmark AUC requires inserting a query token. " + "Please evaluate only runs trained with the current no-event vocabulary." + ) + return + + subset_indices = make_eval_indices(dataset, args, cfg) + + first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps( + dataset, subset_indices) + + 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=first_occurrence_by_token, + ) + + if not disease_ids: + raise RuntimeError("No disease tokens selected after filtering.") + + landmark_start = float(cfg_get(args, cfg, "landmark_start", 40.0)) + landmark_stop = float(cfg_get(args, cfg, "landmark_stop", 80.0)) + landmark_step = float(cfg_get(args, cfg, "landmark_step", 5.0)) + if landmark_step <= 0: + raise ValueError("landmark_step must be > 0") + landmark_ages = np.arange( + landmark_start, landmark_stop, landmark_step, dtype=np.float32) + if landmark_ages.size == 0: + raise ValueError( + "Landmark ages are empty. Check landmark_start/landmark_stop/landmark_step.") + + horizons = np.asarray( + parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [ + 1.0, 5.0, 10.0], + dtype=np.float32, + ) + if horizons.size == 0: + raise ValueError("horizons must contain at least one value") + + score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower() + if score_mode not in {"risk", "eta"}: + raise ValueError("score_mode must be 'risk' or 'eta'") + + state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu") + dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict) + if dist_mode == "weibull": + raise RuntimeError( + "Weibull checkpoints are not supported by evaluate_auc_v2.py. " + "This landmark evaluation requires imputation, and Weibull runs " + "in this project are treated as unsupported for no-event landmark AUC. " + "Please use an exponential no-event checkpoint." + ) + if score_mode == "risk" and dist_mode != "exponential": + raise RuntimeError( + "score_mode='risk' requires dist_mode='exponential' for evaluate_auc_v2.py." + ) + + if score_mode == "eta": + print("WARNING: eta diagnostic score is not horizon-specific risk.") + + cfg_model = dict(cfg) + cfg_model["dist_mode"] = dist_mode + + device = torch.device(cfg.get("device", "cuda") + if torch.cuda.is_available() else "cpu") + if device.type == "cuda": + torch.backends.cudnn.benchmark = True + + model = build_model_from_dataset(args, cfg_model, dataset).to(device) + + if model.token_embedding.num_embeddings <= NO_EVENT_IDX or model.risk_head.out_features <= NO_EVENT_IDX: + raise RuntimeError( + "Model vocabulary does not include token index. " + "Checkpoint/model shape is incompatible with no-event landmark querying." + ) + + try: + load_model_state(model, state_dict) + except RuntimeError as exc: + raise RuntimeError( + "Checkpoint vocabulary shape is incompatible with the no-event dataset/model setup. " + "Please ensure this run was trained with the current no-event vocabulary and matching labels file." + ) from exc + + if model.token_embedding.num_embeddings != dataset.vocab_size or model.risk_head.out_features != dataset.vocab_size: + raise RuntimeError( + "Checkpoint/model vocabulary shape mismatch with dataset vocabulary. " + "Please use a checkpoint trained with the same no-event vocabulary configuration." + ) + + death_token_ids = _get_death_token_ids(dataset, labels_meta) + min_history_events = int(cfg_get(args, cfg, "min_history_events", 1)) + landmark_dataset = LandmarkDataset( + dataset=dataset, + subset_indices=subset_indices, + landmark_ages=landmark_ages, + attn_mask_mode=attn_mask_mode, + min_history_events=min_history_events, + first_occurrence_by_token=first_occurrence_by_token, + death_token_ids=death_token_ids, + ) + + batch_size = int(cfg_get(args, cfg, "batch_size", 128)) + num_workers = int(cfg_get(args, cfg, "num_workers", 4)) + loader = DataLoader( + landmark_dataset, + batch_size=batch_size, + shuffle=False, + collate_fn=collate_landmark_fn, + num_workers=num_workers, + pin_memory=device.type == "cuda", + persistent_workers=num_workers > 0, + prefetch_factor=2 if num_workers > 0 else None, + ) + + eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower() + if eval_split in {"valid", "validation"}: + eval_split = "val" + + score_mode_out = ( + "insert_no_event_landmark_exponential_risk" + if score_mode == "risk" + else "insert_no_event_landmark_eta_diagnostic" + ) + + 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)) + disease_chunk_size = int(cfg_get(args, cfg, "disease_chunk_size", 0)) + min_cases = int(cfg_get(args, cfg, "min_cases", 2)) + exclude_death_competing = bool( + cfg_get(args, cfg, "exclude_death_in_window_without_disease", True)) + use_amp = bool(cfg_get(args, cfg, "use_amp", False)) + hidden_cache_dtype = str( + cfg_get(args, cfg, "hidden_cache_dtype", "float16")) + logit_batch_size = int(cfg_get(args, cfg, "logit_batch_size", batch_size)) + + print(f"Eval split: {eval_split}") + print(f"Number of selected patients: {len(subset_indices)}") + print("No-event support: true") + print("Landmark query mode: insert_no_event_token") + print("Landmark token mode: no_event") + print(f"Dist mode: {dist_mode}") + print(f"Score mode: {score_mode}") + print(f"Landmark ages: {landmark_ages.tolist()}") + print(f"Horizons: {horizons.tolist()}") + print(f"Number of landmark query samples: {len(landmark_dataset)}") + print(f"Number of disease tokens: {len(disease_ids)}") + print(f"AUC workers: {num_workers_auc}") + print(f"Output path: {output_path}") + + meta_info = { + "score_mode": score_mode_out, + "eval_split": eval_split, + "model_ckpt_path": str(model_ckpt_path), + "config_path": str(config_path), + "target_mode": str(target_mode), + "dist_mode": str(dist_mode), + "time_mode": str(time_mode), + "attn_mask_mode": str(attn_mask_mode), + "readout_name": str(readout_name), + "landmark_query_mode": "insert_no_event_token", + "landmark_token_mode": "no_event", + } + + evaluate_landmark_auc( + model=model, + loader=loader, + landmark_dataset=landmark_dataset, + output_path=output_path, + labels_meta=labels_meta, + disease_ids=disease_ids, + disease_chunk_size=disease_chunk_size, + score_mode=score_mode, + horizons=horizons, + device=device, + attn_mask_mode=attn_mask_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + num_workers_auc=num_workers_auc, + auc_task_chunk_size=auc_task_chunk_size, + min_cases=min_cases, + exclude_death_competing=exclude_death_competing, + use_amp=use_amp, + hidden_cache_dtype=hidden_cache_dtype, + logit_batch_size=logit_batch_size, + meta_info=meta_info, + ) + + +if __name__ == "__main__": + main()