"""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()