"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth. This script supports DeepHealth fixed-horizon risk scores for exponential and Weibull all-future distributions. The default horizons are 0.1, 1, 5, and 10 years. As in Delphi2M, 0.1 years is reported as the no-gap evaluation. Landmark querying depends on the model target mode saved in train_config.json: - next_token: insert a token at landmark age and read it out; - all_future: pass landmark age directly as t_query. """ 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 import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence from torch.utils.data import DataLoader, Dataset from tqdm.auto import tqdm from dataset import ( DISEASE_HISTORY_MODE_TIMED, HealthDataset, normalize_disease_history_mode, transform_disease_history, ) from delphi2m_auc_report import ( DEFAULT_DELPHI2M_PERIODS_YEARS, build_delphi2m_auc_report, ) from eval_data import ( build_first_occurrence_map, build_model_from_dataset, cfg_get, load_json_config, load_sequence_eval_dataset, resolve_eval_device, split_indices, validate_training_mode_config, validate_dataset_metadata, ) from model_architectures import resolve_model_architecture from models import DeepHealth from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX} 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 if text.startswith("["): try: values = json.loads(text) except json.JSONDecodeError as exc: raise ValueError(f"Invalid integer list: {text!r}") from exc if not isinstance(values, list): raise ValueError(f"Expected a JSON list, got {type(values).__name__}") return [int(x) for x in values] 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 if text.startswith("["): try: values = json.loads(text) except json.JSONDecodeError as exc: raise ValueError(f"Invalid float list: {text!r}") from exc if not isinstance(values, list): raise ValueError(f"Expected a JSON list, got {type(values).__name__}") return [float(x) for x in values] return [float(x.strip()) for x in text.split(",") if x.strip()] 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() if mode not in {"exponential", "weibull"}: raise ValueError( f"Unsupported dist_mode={mode!r}; expected exponential or weibull." ) has_rho_head = any(str(k).startswith("rho_head.") for k in state_dict.keys()) if mode == "weibull" and not has_rho_head: raise RuntimeError("Weibull checkpoint is missing rho_head parameters.") if mode == "exponential" and has_rho_head: raise RuntimeError( "Exponential checkpoint unexpectedly contains rho_head parameters." ) return mode def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None: resolve_model_architecture(model.model_architecture, state_dict) model.load_state_dict(state_dict, strict=True) # --------------------------------------------------------------------------- # 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 _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) -> List[int]: return [int(dataset.vocab_size) - 1] 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, model_target_mode: str, min_history_events: int, first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], death_token_ids: Sequence[int], disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED, ) -> 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.model_target_mode = str(model_target_mode).lower() if self.model_target_mode not in {"next_token", "all_future"}: raise ValueError( "model_target_mode must be next_token or all_future, got " f"{self.model_target_mode!r}" ) self.disease_history_mode = normalize_disease_history_mode( disease_history_mode ) if ( self.model_target_mode != "all_future" and self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED ): raise ValueError( "ordered/set disease history is only supported for all_future models" ) 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, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64)) if valid_history_mask.sum() < self.min_history_events: continue if self.model_target_mode == "next_token": 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), ] ) time_seq_landmark[-1] = np.nextafter( np.float32(landmark_age), np.float32(np.inf), dtype=np.float32, ) landmark_pos = int(len(event_seq_landmark) - 1) readout_mask = np.zeros(len(event_seq_landmark), dtype=bool) readout_mask[-1] = True else: ( event_seq_landmark, time_seq_landmark, model_t_query, ) = transform_disease_history( event_seq=prefix_events, actual_time_seq=prefix_times, actual_t_query=landmark_age, disease_history_mode=self.disease_history_mode, ) landmark_pos = int(len(event_seq_landmark) - 1) readout_mask = np.zeros(len(event_seq_landmark), dtype=bool) 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": landmark_pos, "t_query": ( np.float32(landmark_age) if self.model_target_mode == "next_token" else model_t_query ), "event_seq": event_seq_landmark, "time_seq": time_seq_landmark, "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(), "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), "t_query": torch.tensor(float(s["t_query"]), dtype=torch.float32), "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) 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, "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]), "t_query": torch.stack([x["t_query"] 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, model_target_mode: str, use_amp: bool, hidden_cache_dtype: str, ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: model_target_mode = str(model_target_mode).lower() if model_target_mode not in {"next_token", "all_future"}: raise ValueError( f"model_target_mode must be next_token or all_future, got {model_target_mode!r}" ) 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: if model_target_mode == "all_future": landmark_hidden = model( event_seq=batch_dev["event_seq"], time_seq=batch_dev["time_seq"], sex=batch_dev["sex"], padding_mask=batch_dev["padding_mask"], t_query=batch_dev["t_query"], other_type=batch_dev["other_type"], other_value=batch_dev["other_value"], other_value_kind=batch_dev["other_value_kind"], other_time=batch_dev["other_time"], ) else: 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"], ) landmark_hidden = hidden.gather( 1, batch_dev["landmark_pos"].long()[:, None, None].expand( -1, 1, 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_distribution_chunk( model: DeepHealth, hidden_all: np.ndarray, disease_ids: Sequence[int], dist_mode: str, device: torch.device, logit_batch_size: int, use_amp: bool, ) -> Tuple[np.ndarray, Optional[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] dist_mode = str(dist_mode).lower() 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) bias = None if model.risk_head.bias is not None: bias = model.risk_head.bias[disease_ids].detach().to( device=device, dtype=compute_dtype) rho_weight = None rho_bias = None if dist_mode == "weibull": rho_weight = model.rho_head.weight[disease_ids].detach().to( device=device, dtype=compute_dtype) rho_bias = model.rho_head.bias[disease_ids].detach().to( device=device, dtype=compute_dtype) out_parts: List[np.ndarray] = [] rho_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()) if bias is not None: logits = logits + bias rho = None if dist_mode == "weibull": assert rho_weight is not None and rho_bias is not None rho = F.softplus(torch.matmul(h, rho_weight.t()) + rho_bias) + 1e-6 out_parts.append(logits.float().cpu( ).numpy().astype(np.float32, copy=False)) if rho is not None: rho_parts.append(rho.float().cpu( ).numpy().astype(np.float32, copy=False)) del h, logits, rho logits_all = np.concatenate(out_parts, axis=0) rho_all = np.concatenate(rho_parts, axis=0) if rho_parts else None return logits_all, rho_all # --------------------------------------------------------------------------- # Parallel AUC workers # --------------------------------------------------------------------------- _WORKER: Dict[str, Any] = {} def _init_worker( disease_ids: np.ndarray, score_chunk: np.ndarray, rho_chunk: Optional[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, dist_mode: str, ) -> 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), "rho_chunk": None if rho_chunk is None else np.asarray(rho_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()), "dist_mode": str(dist_mode).lower(), "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, rho: Optional[np.ndarray], score_mode: str, horizon: float, dist_mode: str, ) -> 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) dist_mode = str(dist_mode).lower() if dist_mode == "weibull": if rho is None: raise RuntimeError("Weibull risk scoring requires rho parameters.") exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False)) return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False) 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)] rho_chunk = _WORKER["rho_chunk"] rho_token = None if rho_chunk is None else rho_chunk[:, int(j)] dist_mode = _WORKER["dist_mode"] 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]], None if rho_token is None else rho_token[idx[case_idx]], score_mode=score_mode, horizon=horizon, dist_mode=dist_mode, ) control_scores = _score_to_probability( logits_token[idx[control_idx]], None if rho_token is None else rho_token[idx[control_idx]], score_mode=score_mode, horizon=horizon, dist_mode=dist_mode, ) 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, disease_ids: Sequence[int], disease_chunk_size: int, score_mode: str, dist_mode: str, horizons: np.ndarray, device: torch.device, model_target_mode: 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, ) -> Tuple[pd.DataFrame, pd.DataFrame]: model.eval().to(device) hidden_all, row_arrays = infer_landmark_hidden( model=model, loader=loader, device=device, model_target_mode=model_target_mode, 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, rho_chunk = project_distribution_chunk( model=model, hidden_all=hidden_all, disease_ids=chunk_ids, dist_mode=dist_mode, 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, rho_chunk=rho_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), dist_mode=dist_mode, ) 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, rho_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), dist_mode, ), ) 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, rho_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) print( "Building Delphi2M-style report: mean AUC across landmark-age " "strata, reported separately for Female and Male." ) df_report = build_delphi2m_auc_report( df_unpooled, period_col="horizon", ) output_path.mkdir(parents=True, exist_ok=True) report_path = output_path / "df_auc_landmark_delphi2m_report.csv" df_report.to_csv(report_path, index=False) print(f"Saved Delphi2M-style landmark AUC report: {report_path}") return df_unpooled, df_report 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("--device", type=str, default=None, help="Evaluation device, e.g. cpu, cuda, cuda:1. Defaults to cuda if available, else cpu.") 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, help="Comma-separated horizons in years; defaults to 0.1,1,5,10, where 0.1 is Delphi2M no gap.", ) 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) validate_training_mode_config(cfg) 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) model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() if model_target_mode not in {"next_token", "all_future"}: raise ValueError( "train_config.json model_target_mode must be next_token or all_future, " f"got {model_target_mode!r}" ) dist_mode_cfg = str(cfg.get("dist_mode", "exponential")) disease_history_mode = normalize_disease_history_mode( cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED) ) 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 = load_sequence_eval_dataset( model_target_mode=model_target_mode, data_prefix=data_prefix, labels_file=labels_file, no_event_interval_years=float(no_event_interval_years), min_history_events=int(cfg.get("all_future_min_history_events", 1)), min_future_events=int(cfg.get("all_future_min_future_events", 1)), extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), disease_history_mode=disease_history_mode, ) validate_dataset_metadata(dataset, cfg) 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 model_target_mode == "next_token" and 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_map( 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", "0.1,1,5,10") ) or list(DEFAULT_DELPHI2M_PERIODS_YEARS), 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 score_mode == "eta": print( "WARNING: eta diagnostic score is not horizon-specific risk and does not use dist_mode-specific rho parameters.") cfg_model = dict(cfg) cfg_model["dist_mode"] = dist_mode model_architecture = resolve_model_architecture(cfg_model, state_dict) cfg_model["model_architecture"] = model_architecture print(f"Resolved model architecture: {model_architecture}") device = resolve_eval_device(args.device) if device.type == "cuda": torch.backends.cudnn.benchmark = True model = build_model_from_dataset( args, cfg_model, dataset, state_dict=state_dict ).to(device) if ( model_target_mode == "next_token" and ( model.token_embedding.num_embeddings <= NO_EVENT_IDX or model.risk_head.out_features <= NO_EVENT_IDX ) ): raise RuntimeError( "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 dataset/model setup. " "Please ensure this run was trained with the same special-token vocabulary and 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) 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, model_target_mode=model_target_mode, min_history_events=min_history_events, first_occurrence_by_token=first_occurrence_by_token, death_token_ids=death_token_ids, disease_history_mode=disease_history_mode, ) 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" landmark_query_mode = ( "insert_no_event_token" if model_target_mode == "next_token" else "direct_t_query" ) 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(f"No-event support: {bool(has_no_event)}") print(f"Model target mode: {model_target_mode}") print(f"Disease history mode: {disease_history_mode}") print(f"Landmark query mode: {landmark_query_mode}") print( "Landmark token mode: no_event" if model_target_mode == "next_token" else "Landmark token mode: none" ) 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}") evaluate_landmark_auc( model=model, loader=loader, landmark_dataset=landmark_dataset, output_path=output_path, disease_ids=disease_ids, disease_chunk_size=disease_chunk_size, score_mode=score_mode, dist_mode=dist_mode, horizons=horizons, device=device, model_target_mode=model_target_mode, 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, ) if __name__ == "__main__": main()