"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth. This script supports DeepHealth fixed-horizon risk scores for exponential, Weibull, and mixed all-future distributions. 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 HealthDataset from eval_data import load_sequence_eval_dataset 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 resolve_eval_device(device_arg: Optional[str]) -> torch.device: """Resolve evaluation device without inheriting train_config.json device.""" device_name = device_arg or ("cuda" if torch.cuda.is_available() else "cpu") device = torch.device(device_name) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError( f"Requested device {device_name!r}, but CUDA is not available." ) return device 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 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()) has_rho_death_head = any(str(k).startswith("rho_death_head.") for k in state_dict.keys()) if has_rho_head: if mode != "weibull": print( "[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.") return "weibull" if has_rho_death_head: if mode != "mixed": print( "[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.") return "mixed" if mode == "weibull": print( "[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.") return "exponential" if mode == "mixed": print( "[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.") return "exponential" return mode if mode in {"exponential", "weibull", "mixed"} else "exponential" def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth: model_target_mode = str(cfg_get( args, cfg, "model_target_mode", "next_token")).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}" ) 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)), extra_pool_reduce=str(cfg_get(args, cfg, "extra_pool_reduce", "mean")), target_mode=model_target_mode, 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: model.load_state_dict(state_dict, strict=True) def validate_dataset_metadata(dataset: HealthDataset, cfg: Dict[str, Any]) -> None: meta = cfg.get("dataset_metadata") if not isinstance(meta, dict): return actual: Dict[str, Any] = { "vocab_size": int(dataset.vocab_size), "n_types": int(dataset.n_types), "n_cont_types": int(dataset.n_cont_types), "n_categories": int(dataset.n_categories), "cont_type_ids": [int(x) for x in dataset.cont_type_ids], "extra_info_types": [int(x) for x in dataset.extra_info_types], } mismatches = [ f"{key}: train_config={meta.get(key)!r}, current_dataset={value!r}" for key, value in actual.items() if key in meta and meta.get(key) != value ] if mismatches: raise RuntimeError( "Current dataset metadata does not match train_config.json. " "Use the same prepared data and extra_info_types as training. " + "; ".join(mismatches) ) # --------------------------------------------------------------------------- # 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, model_target_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.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.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 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), ] ) if self.attn_mask_mode in _TARGET_AWARE_MODES: 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 = prefix_events.astype( np.int64, copy=False) time_seq_landmark = prefix_times.astype( np.float32, copy=False) 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), "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, readout_name: str, readout_reduce: 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}" ) readout = None if model_target_mode == "next_token" and readout_name == "same_time_group_end": readout = build_readout("same_time_group_end", reduce=readout_reduce).to(device) elif model_target_mode == "next_token": readout = build_readout(readout_name).to(device) if readout is not None: 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: 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"], target_mode="all_future", ) 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"], 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_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 death_rho_weight = None death_rho_bias = None mixed_death_cols: List[int] = [] death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1)) 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) elif dist_mode == "mixed": mixed_death_cols = [j for j, token in enumerate(disease_ids) if int(token) == death_idx] if mixed_death_cols: death_rho_weight = model.rho_death_head.weight.detach().to( device=device, dtype=compute_dtype) death_rho_bias = model.rho_death_head.bias.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 elif dist_mode == "mixed" and mixed_death_cols: assert death_rho_weight is not None and death_rho_bias is not None rho = torch.ones_like(logits) death_rho = F.softplus( torch.matmul(h, death_rho_weight.t()).squeeze(-1) + death_rho_bias.squeeze(0) ) + 1e-6 for col in mixed_death_cols: rho[:, int(col)] = death_rho 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, model_death_idx: int, ) -> None: os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("MKL_NUM_THREADS", "1") os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") os.environ.setdefault("NUMEXPR_NUM_THREADS", "1") _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(), "model_death_idx": int(model_death_idx), "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, token: int, death_idx: int, ) -> 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) if dist_mode == "mixed" and int(token) == int(death_idx): if rho is None: raise RuntimeError("Mixed death risk scoring requires death 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"] model_death_idx = int(_WORKER["model_death_idx"]) 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, token=token, death_idx=model_death_idx, ) 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, token=token, death_idx=model_death_idx, ) 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, dist_mode: str, horizons: np.ndarray, device: torch.device, model_target_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, model_target_mode=model_target_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, 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, model_death_idx=int(getattr( model, "death_idx", getattr(model, "vocab_size", 1) - 1)), ) 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, int(getattr( model, "death_idx", getattr(model, "vocab_size", 1) - 1)), ), ) 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) 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", "model_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("--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) 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") 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")) 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 = 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), include_no_event_in_uts_target=bool(include_no_event_in_uts_target), 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)), ) 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_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 not in {"exponential", "weibull", "mixed"}: raise ValueError( f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed." ) 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 device = resolve_eval_device(args.device) if device.type == "cuda": torch.backends.cudnn.benchmark = True model = build_model_from_dataset(args, cfg_model, dataset).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, 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, 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, ) 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" ) score_mode_out = f"{landmark_query_mode}_{score_mode}" 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"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}") 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), "model_target_mode": str(model_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": landmark_query_mode, "landmark_token_mode": "no_event" if model_target_mode == "next_token" else "none", } 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, dist_mode=dist_mode, horizons=horizons, device=device, model_target_mode=model_target_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()