1424 lines
57 KiB
Python
1424 lines
57 KiB
Python
"""
|
|
Evaluate disease-specific AUC for DeepHealth models.
|
|
|
|
This script follows the logic of the Delphi evaluation script supplied by the user:
|
|
1. choose disease tokens of interest from dataset vocabulary;
|
|
2. precompute, for each patient/target occurrence, the latest prediction token
|
|
at least `offset` years before the target time;
|
|
3. run model inference by disease chunks to avoid materializing all logits;
|
|
4. compute AUC separately by sex and age bracket;
|
|
5. aggregate age brackets with DeLong variance.
|
|
|
|
Efficiency notes:
|
|
- transformer/readout inference is executed once and cached;
|
|
- disease chunks reuse the cached hidden states and only recompute selected risk-head columns;
|
|
- AUC work is parallelized on CPU across disease task blocks using process workers;
|
|
- per-sex data are compacted into an event-level table before multiprocessing;
|
|
- large per-sex arrays are installed once per worker with fork-style globals on Linux,
|
|
avoiding repeated pickling of arrays for every disease.
|
|
|
|
Run from the DeepHealth code directory containing dataset.py, models.py,
|
|
readouts.py, and train_config.json-compatible checkpoints/configs.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import math
|
|
import multiprocessing as mp
|
|
import os
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from torch.utils.data import DataLoader, Subset
|
|
from tqdm.auto import tqdm
|
|
|
|
from dataset import HealthDataset
|
|
from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
|
|
from models import DeepHealth
|
|
from readouts import build_readout
|
|
from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeLong AUC utilities, adapted from the supplied Delphi evaluation script
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def compute_midrank(x: np.ndarray) -> np.ndarray:
|
|
"""Compute midranks for DeLong variance."""
|
|
x = np.asarray(x, dtype=np.float64)
|
|
order = np.argsort(x)
|
|
sorted_x = x[order]
|
|
n = len(x)
|
|
ranks = np.zeros(n, dtype=np.float64)
|
|
i = 0
|
|
while i < n:
|
|
j = i
|
|
while j < n and sorted_x[j] == sorted_x[i]:
|
|
j += 1
|
|
ranks[i:j] = 0.5 * (i + j - 1)
|
|
i = j
|
|
out = np.empty(n, dtype=np.float64)
|
|
out[order] = ranks + 1.0
|
|
return out
|
|
|
|
|
|
def fast_delong(predictions_sorted_transposed: np.ndarray, label_1_count: int) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""
|
|
Fast DeLong covariance for ROC AUC.
|
|
|
|
This evaluation uses one classifier at a time, so k is normally 1.
|
|
In that case np.cov(v01) and np.cov(v10) must be scalar variances.
|
|
Do NOT expand them to (m,m)/(n,n) covariance matrices; that is the
|
|
source of the broadcast error:
|
|
shapes (n_case,n_case) and (n_control,n_control)
|
|
"""
|
|
predictions_sorted_transposed = np.asarray(
|
|
predictions_sorted_transposed, dtype=np.float64)
|
|
m = int(label_1_count)
|
|
n = int(predictions_sorted_transposed.shape[1] - m)
|
|
if m <= 0 or n <= 0:
|
|
return np.array([np.nan], dtype=np.float64), np.array([[np.nan]], dtype=np.float64)
|
|
|
|
positive_examples = predictions_sorted_transposed[:, :m]
|
|
negative_examples = predictions_sorted_transposed[:, m:]
|
|
k = int(predictions_sorted_transposed.shape[0])
|
|
|
|
tx = np.empty((k, m), dtype=np.float64)
|
|
ty = np.empty((k, n), dtype=np.float64)
|
|
tz = np.empty((k, m + n), dtype=np.float64)
|
|
for r in range(k):
|
|
tx[r] = compute_midrank(positive_examples[r])
|
|
ty[r] = compute_midrank(negative_examples[r])
|
|
tz[r] = compute_midrank(predictions_sorted_transposed[r])
|
|
|
|
aucs = tz[:, :m].sum(axis=1) / m / n - float(m + 1.0) / 2.0 / n
|
|
|
|
# DeLong structural components:
|
|
# v01: classifier x cases
|
|
# v10: classifier x controls
|
|
v01 = (tz[:, :m] - tx) / n
|
|
v10 = 1.0 - (tz[:, m:] - ty) / m
|
|
|
|
if k == 1:
|
|
# np.cov on a single row is easy to misuse. The correct covariance
|
|
# for one classifier is simply the scalar sample variance over cases
|
|
# plus the scalar sample variance over controls.
|
|
sx = np.var(v01[0], ddof=1) if m > 1 else 0.0
|
|
sy = np.var(v10[0], ddof=1) if n > 1 else 0.0
|
|
delong_cov = np.array([[sx / m + sy / n]], dtype=np.float64)
|
|
else:
|
|
# Multiple-classifier general case. rowvar=True: rows are classifiers.
|
|
sx = np.cov(v01, rowvar=True) if m > 1 else np.zeros(
|
|
(k, k), dtype=np.float64)
|
|
sy = np.cov(v10, rowvar=True) if n > 1 else np.zeros(
|
|
(k, k), dtype=np.float64)
|
|
delong_cov = np.atleast_2d(sx) / m + np.atleast_2d(sy) / n
|
|
|
|
return aucs, delong_cov
|
|
|
|
|
|
def get_auc_delong_var(control_scores: np.ndarray, case_scores: np.ndarray) -> Tuple[float, float]:
|
|
"""Return AUC and DeLong variance for controls/class-0 and cases/class-1."""
|
|
control_scores = np.asarray(control_scores, dtype=np.float64)
|
|
case_scores = np.asarray(case_scores, dtype=np.float64)
|
|
if len(control_scores) == 0 or len(case_scores) == 0:
|
|
return np.nan, np.nan
|
|
|
|
ground_truth = np.array([1] * len(case_scores) +
|
|
[0] * len(control_scores), dtype=np.int8)
|
|
predictions = np.concatenate(
|
|
[case_scores, control_scores]).astype(np.float64, copy=False)
|
|
order = (-ground_truth).argsort()
|
|
label_1_count = int(ground_truth.sum())
|
|
|
|
aucs, delong_cov = fast_delong(
|
|
predictions[np.newaxis, order], label_1_count)
|
|
var = float(np.asarray(delong_cov).reshape(-1)[0])
|
|
return float(aucs[0]), var
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Disease selection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
|
|
|
|
|
|
def _build_first_occurrence_maps(
|
|
dataset: HealthDataset,
|
|
subset_indices: np.ndarray,
|
|
) -> Tuple[Dict[int, Tuple[np.ndarray, np.ndarray]], np.ndarray, np.ndarray, np.ndarray]:
|
|
patient_count = len(subset_indices)
|
|
followup_end = np.full(patient_count, -np.inf, dtype=np.float32)
|
|
death_time = np.full(patient_count, np.inf, dtype=np.float32)
|
|
sex = np.full(patient_count, -1, dtype=np.int8)
|
|
|
|
first_lists: Dict[int, List[Tuple[int, float]]] = {}
|
|
|
|
for patient_id, dataset_index in enumerate(subset_indices.tolist()):
|
|
s = dataset.samples[int(dataset_index)]
|
|
|
|
seq_event = np.asarray(s["event_seq"], dtype=np.int64)
|
|
seq_time = np.asarray(s["time_seq"], dtype=np.float32)
|
|
tgt_event = np.asarray(s["target_event_seq"], dtype=np.int64)
|
|
tgt_time = np.asarray(s["target_time_seq"], dtype=np.float32)
|
|
|
|
if seq_event.size == 0 or tgt_event.size == 0:
|
|
continue
|
|
|
|
full_event = np.concatenate([seq_event, tgt_event[-1:]])
|
|
full_time = np.concatenate([seq_time, tgt_time[-1:]])
|
|
|
|
sex[patient_id] = int(s["sex"])
|
|
followup_end[patient_id] = np.max(full_time).astype(np.float32)
|
|
|
|
uniq_tokens, first_idx = np.unique(full_event, return_index=True)
|
|
|
|
for token, idx in zip(uniq_tokens.tolist(), first_idx.tolist()):
|
|
token = int(token)
|
|
event_time = float(full_time[int(idx)])
|
|
if token not in first_lists:
|
|
first_lists[token] = []
|
|
first_lists[token].append((patient_id, event_time))
|
|
|
|
packed: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
|
|
for token, pairs in first_lists.items():
|
|
if not pairs:
|
|
continue
|
|
packed[int(token)] = (
|
|
np.asarray([p for p, _ in pairs], dtype=np.int32),
|
|
np.asarray([t for _, t in pairs], dtype=np.float32),
|
|
)
|
|
|
|
return packed, followup_end, death_time, sex
|
|
|
|
|
|
def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
|
|
death_ids = [int(dataset.vocab_size) - 1]
|
|
print(f"[INFO] death token ids: {death_ids}")
|
|
return death_ids
|
|
|
|
|
|
def select_disease_tokens(
|
|
dataset: HealthDataset,
|
|
requested_tokens: Optional[Sequence[int]],
|
|
filter_min_total: int,
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
) -> List[int]:
|
|
base = [
|
|
int(token)
|
|
for token, code in dataset.label_id_to_code.items()
|
|
if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<")
|
|
]
|
|
base_set = set(base)
|
|
print(f"[INFO] Valid disease tokens in current vocabulary: {len(base)}")
|
|
|
|
if requested_tokens is not None:
|
|
selected = sorted(
|
|
set(int(x) for x in requested_tokens if int(x) in base_set and int(x) not in SPECIAL_TOKENS))
|
|
print(
|
|
"[INFO] Requested disease tokens provided: "
|
|
f"input={len(list(requested_tokens))}, selected={len(selected)} (vocab/SPECIAL filtered).")
|
|
print(f"[INFO] Final disease_ids count: {len(selected)}")
|
|
return selected
|
|
|
|
disease_ids = sorted(base)
|
|
if int(filter_min_total) <= 0:
|
|
print(
|
|
f"[INFO] filter_min_total={int(filter_min_total)} <= 0; keeping all {len(disease_ids)} disease tokens.")
|
|
print(f"[INFO] Final disease_ids count: {len(disease_ids)}")
|
|
return disease_ids
|
|
|
|
print(
|
|
f"[INFO] Applying filter_min_total={int(filter_min_total)}: before={len(disease_ids)} tokens.")
|
|
print(
|
|
"[INFO] Using split first-occurrence patient counts for filter_min_total.")
|
|
split_counts = {}
|
|
for token, pairs in first_occurrence_by_token.items():
|
|
token = int(token)
|
|
if token not in base_set:
|
|
continue
|
|
split_counts[token] = len(np.unique(pairs[0]))
|
|
filtered = [token for token in disease_ids if int(
|
|
split_counts.get(token, 0)) > int(filter_min_total)]
|
|
print(
|
|
"[INFO] First-occurrence count filtering complete: "
|
|
f"after={len(filtered)} tokens.")
|
|
print(f"[INFO] Final disease_ids count: {len(filtered)}")
|
|
return filtered
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dataset/split/model helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_json_config(path: Optional[str]) -> Dict[str, Any]:
|
|
if path is None:
|
|
return {}
|
|
p = Path(path)
|
|
if not p.exists():
|
|
return {}
|
|
with p.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def cfg_get(args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any], name: str, default: Any) -> Any:
|
|
"""Get a value from CLI args first, then train_config.json, then default.
|
|
|
|
This helper intentionally accepts either an argparse.Namespace or a dict.
|
|
The earlier version passed cfg as both args and cfg, then tried to access
|
|
args.eval_split, which fails because dict has no attributes.
|
|
"""
|
|
val = None
|
|
if args is not None:
|
|
if isinstance(args, dict):
|
|
val = args.get(name, None)
|
|
else:
|
|
val = getattr(args, name, None)
|
|
if val is not None:
|
|
return val
|
|
return cfg.get(name, default)
|
|
|
|
|
|
def 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 split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: float, seed: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
total = train_ratio + val_ratio + test_ratio
|
|
if not np.isclose(total, 1.0, atol=1e-6):
|
|
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
|
rng = np.random.RandomState(seed)
|
|
idx = rng.permutation(n)
|
|
n_train = int(n * train_ratio)
|
|
n_val = int(n * val_ratio)
|
|
return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
|
|
|
|
|
|
def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
|
|
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)),
|
|
target_mode=model_target_mode,
|
|
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
|
|
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
|
|
use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)),
|
|
exposure_d_model=cfg_get(args, cfg, "exposure_d_model", None),
|
|
exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
|
|
exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 3)),
|
|
exposure_n_convnext_blocks=int(cfg_get(args, cfg, "exposure_n_convnext_blocks", 2)),
|
|
exposure_conv_kernel_size=int(cfg_get(args, cfg, "exposure_conv_kernel_size", 7)),
|
|
exposure_mlp_ratio=float(cfg_get(args, cfg, "exposure_mlp_ratio", 4.0)),
|
|
exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)),
|
|
)
|
|
|
|
|
|
def _extract_state_dict(ckpt: Any) -> Dict[str, Any]:
|
|
if isinstance(ckpt, dict) and "model" in ckpt:
|
|
return ckpt["model"]
|
|
elif isinstance(ckpt, dict) and "state_dict" in ckpt:
|
|
return ckpt["state_dict"]
|
|
return ckpt
|
|
|
|
|
|
def load_checkpoint_state_dict(checkpoint_path: str, map_location: str | torch.device = "cpu") -> Dict[str, Any]:
|
|
ckpt = torch.load(checkpoint_path, map_location=map_location)
|
|
state = _extract_state_dict(ckpt)
|
|
if not isinstance(state, dict):
|
|
raise TypeError(
|
|
f"Unsupported checkpoint payload type: {type(state)}")
|
|
return state
|
|
|
|
|
|
def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
|
|
mode = str(cfg_dist_mode).lower()
|
|
has_rho_head = any(str(k).startswith("rho_head.")
|
|
for k in state_dict.keys())
|
|
has_rho_death_head = any(str(k).startswith("rho_death_head.")
|
|
for k in state_dict.keys())
|
|
|
|
if has_rho_head and mode != "weibull":
|
|
print(
|
|
"[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
|
|
return "weibull"
|
|
if has_rho_death_head and mode != "mixed":
|
|
print(
|
|
"[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
|
|
return "mixed"
|
|
if (not has_rho_head) and mode == "weibull":
|
|
print(
|
|
"[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
|
|
return "exponential"
|
|
if (not has_rho_death_head) and 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
|
|
|
|
|
|
def load_model_state(
|
|
model: torch.nn.Module,
|
|
checkpoint_path: str,
|
|
device: torch.device,
|
|
state_dict: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
state = state_dict if state_dict is not None else load_checkpoint_state_dict(
|
|
checkpoint_path, map_location=device)
|
|
|
|
model.load_state_dict(state, strict=True)
|
|
|
|
|
|
def make_eval_subset(dataset: HealthDataset, args: argparse.Namespace | Dict[str, Any] | None, cfg: Dict[str, Any]) -> Tuple[Subset, np.ndarray]:
|
|
train_ratio = float(cfg_get(args, cfg, "train_ratio", 0.7))
|
|
val_ratio = float(cfg_get(args, cfg, "val_ratio", 0.15))
|
|
test_ratio = float(cfg_get(args, cfg, "test_ratio", 0.15))
|
|
seed = int(cfg_get(args, cfg, "seed", 42))
|
|
eval_split = str(cfg_get(args, cfg, "eval_split", "test")).lower()
|
|
dataset_subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
|
|
|
train_idx, val_idx, test_idx = split_indices(
|
|
len(dataset), train_ratio, val_ratio, test_ratio, seed)
|
|
split_map = {
|
|
"train": train_idx,
|
|
"val": val_idx,
|
|
"valid": val_idx,
|
|
"validation": val_idx,
|
|
"test": test_idx,
|
|
"all": np.arange(len(dataset)),
|
|
}
|
|
if eval_split not in split_map:
|
|
raise ValueError(
|
|
f"eval_split must be one of {sorted(split_map)}, got {eval_split!r}")
|
|
|
|
indices = split_map[eval_split]
|
|
if dataset_subset_size is not None and int(dataset_subset_size) > 0:
|
|
indices = indices[: int(dataset_subset_size)]
|
|
return Subset(dataset, indices.tolist()), np.asarray(indices, dtype=np.int64)
|
|
|
|
|
|
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),
|
|
}
|
|
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 as training. "
|
|
+ "; ".join(mismatches)
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Batched inference + cached hidden states
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _numpy_hidden_dtype(name: str) -> np.dtype:
|
|
name = str(name).lower()
|
|
if name in {"float16", "fp16", "half"}:
|
|
return np.float16
|
|
if name in {"bfloat16", "bf16"}:
|
|
# NumPy has limited bfloat16 support; store as fp16 for compact CPU cache.
|
|
return np.float16
|
|
if name in {"float32", "fp32", "single"}:
|
|
return np.float32
|
|
raise ValueError(
|
|
f"hidden_cache_dtype must be float16 or float32, got {name!r}")
|
|
|
|
|
|
@torch.inference_mode()
|
|
def infer_readout_hidden(
|
|
model: DeepHealth,
|
|
loader: DataLoader,
|
|
device: torch.device,
|
|
model_target_mode: str,
|
|
readout_name: str,
|
|
readout_reduce: str,
|
|
use_amp: bool,
|
|
hidden_cache_dtype: str = "float16",
|
|
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
|
|
"""Cache per-position hidden states used by the unchanged AUC logic."""
|
|
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: Dict[str, List[np.ndarray]] = {
|
|
"event_seq": [],
|
|
"time_seq": [],
|
|
"target_event_seq": [],
|
|
"target_time_seq": [],
|
|
"padding_mask": [],
|
|
"readout_mask": [],
|
|
"sex": [],
|
|
}
|
|
max_len = 0
|
|
out_dtype = _numpy_hidden_dtype(hidden_cache_dtype)
|
|
autocast_enabled = bool(use_amp and device.type == "cuda")
|
|
|
|
for batch in tqdm(loader, desc="Model/readout inference", leave=False, dynamic_ncols=True):
|
|
batch_dev = {
|
|
k: (v.to(device, non_blocking=True)
|
|
if isinstance(v, torch.Tensor) else v)
|
|
for k, v in batch.items()
|
|
}
|
|
event_seq = batch_dev["event_seq"]
|
|
time_seq = batch_dev["time_seq"]
|
|
padding_mask = batch_dev["padding_mask"]
|
|
|
|
amp_context = (
|
|
torch.autocast(device_type=device.type, dtype=torch.float16)
|
|
if autocast_enabled else contextlib.nullcontext()
|
|
)
|
|
with amp_context:
|
|
if model_target_mode == "all_future":
|
|
batch_size, seq_len = event_seq.shape
|
|
hidden = torch.zeros(
|
|
batch_size,
|
|
seq_len,
|
|
model.n_embd,
|
|
device=event_seq.device,
|
|
dtype=torch.float32,
|
|
)
|
|
for pos in range(seq_len):
|
|
active = padding_mask[:, pos].bool()
|
|
if not active.any():
|
|
continue
|
|
hidden_pos = model(
|
|
event_seq=event_seq[active],
|
|
time_seq=time_seq[active],
|
|
sex=batch_dev["sex"][active],
|
|
padding_mask=padding_mask[active],
|
|
t_query=time_seq[active, pos],
|
|
exposure_daily=(
|
|
batch_dev["exposure_daily"][active]
|
|
if "exposure_daily" in batch_dev
|
|
else None
|
|
),
|
|
exposure_monthly=(
|
|
batch_dev["exposure_monthly"][active]
|
|
if "exposure_monthly" in batch_dev
|
|
else None
|
|
),
|
|
target_mode="all_future",
|
|
)
|
|
hidden[active, pos, :] = hidden_pos.float()
|
|
readout_mask_np = batch["padding_mask"].cpu().numpy()
|
|
else:
|
|
hidden_raw = model(
|
|
event_seq=event_seq,
|
|
time_seq=time_seq,
|
|
sex=batch_dev["sex"],
|
|
padding_mask=padding_mask,
|
|
exposure_daily=batch_dev.get("exposure_daily"),
|
|
exposure_monthly=batch_dev.get("exposure_monthly"),
|
|
target_mode="next_token",
|
|
)
|
|
ro = readout(
|
|
hidden=hidden_raw,
|
|
time_seq=time_seq,
|
|
padding_mask=padding_mask,
|
|
readout_mask=batch_dev["readout_mask"],
|
|
)
|
|
hidden = ro.hidden
|
|
readout_mask_np = ro.readout_mask.detach().cpu().numpy()
|
|
|
|
h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
|
|
hidden_parts.append(h)
|
|
max_len = max(max_len, h.shape[1])
|
|
|
|
for k in arrays:
|
|
if k == "sex":
|
|
arrays[k].append(
|
|
batch[k].cpu().numpy().astype(np.int8, copy=False))
|
|
else:
|
|
arrays[k].append(batch[k].cpu().numpy())
|
|
arrays["readout_mask"][-1] = readout_mask_np
|
|
|
|
def pad_3d(parts: List[np.ndarray], fill: float = 0.0) -> np.ndarray:
|
|
out = np.full(
|
|
(sum(x.shape[0] for x in parts), max_len, parts[0].shape[2]),
|
|
fill,
|
|
dtype=out_dtype,
|
|
)
|
|
s = 0
|
|
for x in parts:
|
|
out[s:s + x.shape[0], :x.shape[1], :] = x
|
|
s += x.shape[0]
|
|
return out
|
|
|
|
def pad_2d(parts: List[np.ndarray], fill: Any = 0, dtype: Optional[np.dtype] = None) -> np.ndarray:
|
|
dtype = dtype or parts[0].dtype
|
|
out = np.full((sum(x.shape[0]
|
|
for x in parts), max_len), fill, dtype=dtype)
|
|
s = 0
|
|
for x in parts:
|
|
out[s:s + x.shape[0], :x.shape[1]] = x
|
|
s += x.shape[0]
|
|
return out
|
|
|
|
hidden_all = pad_3d(hidden_parts)
|
|
arr_out = {
|
|
"event_seq": pad_2d(arrays["event_seq"], PAD_IDX, np.int64),
|
|
"time_seq": pad_2d(arrays["time_seq"], 0.0, np.float32),
|
|
"target_event_seq": pad_2d(arrays["target_event_seq"], PAD_IDX, np.int64),
|
|
"target_time_seq": pad_2d(arrays["target_time_seq"], 0.0, np.float32),
|
|
"padding_mask": pad_2d(arrays["padding_mask"], False, bool),
|
|
"readout_mask": pad_2d(arrays["readout_mask"], False, bool),
|
|
"sex": np.concatenate(arrays["sex"], axis=0),
|
|
}
|
|
return hidden_all, arr_out
|
|
|
|
|
|
@torch.inference_mode()
|
|
def compute_logits_for_disease_chunk(
|
|
model: DeepHealth,
|
|
hidden_all: np.ndarray,
|
|
disease_ids: Sequence[int],
|
|
device: torch.device,
|
|
logit_batch_size: int,
|
|
use_amp: bool,
|
|
) -> np.ndarray:
|
|
"""Project cached hidden states to only the requested disease columns."""
|
|
n = int(hidden_all.shape[0])
|
|
logit_batch_size = max(1, int(logit_batch_size))
|
|
disease_ids = [int(x) for x in disease_ids]
|
|
|
|
compute_dtype = torch.float16 if (
|
|
device.type == "cuda" and use_amp) else torch.float32
|
|
weight = model.risk_head.weight[disease_ids].detach().to(
|
|
device=device, dtype=compute_dtype)
|
|
bias = None
|
|
if model.risk_head.bias is not None:
|
|
bias = model.risk_head.bias[disease_ids].detach().to(
|
|
device=device, dtype=compute_dtype)
|
|
|
|
parts: List[np.ndarray] = []
|
|
for start in tqdm(range(0, n, logit_batch_size), desc="Risk-head projection", leave=False, dynamic_ncols=True):
|
|
end = min(start + logit_batch_size, n)
|
|
h = torch.from_numpy(hidden_all[start:end]).to(
|
|
device=device, dtype=compute_dtype, non_blocking=True)
|
|
logits = torch.matmul(h, weight.t())
|
|
if bias is not None:
|
|
logits = logits + bias
|
|
parts.append(logits.float().cpu().numpy().astype(
|
|
np.float32, copy=False))
|
|
del h, logits
|
|
return np.concatenate(parts, axis=0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CPU-parallel calibration AUC
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_WORKER: Dict[str, Any] = {}
|
|
|
|
|
|
def _build_flat_eval_table(
|
|
p_sex: np.ndarray,
|
|
time_seq: np.ndarray,
|
|
target_time_seq: np.ndarray,
|
|
target_event_seq: np.ndarray,
|
|
padding_mask: np.ndarray,
|
|
readout_mask: np.ndarray,
|
|
offset: float,
|
|
valid_target_min_id: int,
|
|
age_groups: np.ndarray,
|
|
) -> Dict[str, np.ndarray]:
|
|
"""
|
|
Build a compact event-level table for AUC evaluation.
|
|
|
|
The previous version kept 2D patient x position arrays and each disease task
|
|
repeatedly scanned them. This table keeps only eligible target occurrences:
|
|
- target event is a real disease/event target;
|
|
- a valid readout token exists at least `offset` years before target time;
|
|
- the prediction token itself is valid and readout-valid;
|
|
- predicted age falls inside the requested age brackets.
|
|
|
|
The disease-specific logic is unchanged; this only changes the physical layout
|
|
of the data to make CPU multiprocessing cheaper and more cache-friendly.
|
|
"""
|
|
if len(age_groups) < 2:
|
|
raise ValueError("age_groups must contain at least two values")
|
|
age_start = float(age_groups[0])
|
|
age_step = float(age_groups[1] - age_groups[0])
|
|
n_age = int(len(age_groups))
|
|
|
|
# Raw valid-target table is used only to decide which patients are cases for
|
|
# a disease. This intentionally happens before offset/age filtering, matching
|
|
# the supplied Delphi logic: a patient who ever has disease k is not a control
|
|
# for k, even if that occurrence later lacks a valid prediction point.
|
|
raw_patient_idx, raw_target_idx = np.where(
|
|
target_event_seq > int(valid_target_min_id))
|
|
raw_target_event = target_event_seq[raw_patient_idx, raw_target_idx].astype(
|
|
np.int32, copy=False)
|
|
raw_sort_order = np.argsort(raw_target_event, kind="mergesort")
|
|
raw_sorted_target_event = raw_target_event[raw_sort_order]
|
|
raw_patient_idx = raw_patient_idx.astype(np.int32, copy=False)
|
|
|
|
valid_pred_pos = (
|
|
(time_seq[:, :, None] <= (target_time_seq[:, None, :] - float(offset)))
|
|
& readout_mask[:, :, None]
|
|
)
|
|
pos_index = np.arange(time_seq.shape[1], dtype=np.int32)[None, :, None]
|
|
pred_idx_precompute = np.where(
|
|
valid_pred_pos, pos_index, -1).max(axis=1).astype(np.int32)
|
|
|
|
candidate = (target_event_seq > int(valid_target_min_id)) & (
|
|
pred_idx_precompute >= 0)
|
|
patient_idx, target_idx = np.where(candidate)
|
|
if patient_idx.size == 0:
|
|
return {
|
|
"patient": np.empty(0, dtype=np.int32),
|
|
"target_event": np.empty(0, dtype=np.int32),
|
|
"pred_idx": np.empty(0, dtype=np.int32),
|
|
"age_bin": np.empty(0, dtype=np.int16),
|
|
"target_time": np.empty(0, dtype=np.float32),
|
|
"sort_order": np.empty(0, dtype=np.int64),
|
|
"sorted_target_event": np.empty(0, dtype=np.int32),
|
|
"raw_patient": raw_patient_idx,
|
|
"raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
|
|
"raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
|
|
"p_sex": p_sex,
|
|
"age_groups": age_groups.astype(np.float32, copy=False),
|
|
"n_patients": np.int32(time_seq.shape[0]),
|
|
}
|
|
|
|
pred_idx = pred_idx_precompute[patient_idx, target_idx]
|
|
pred_ok = padding_mask[patient_idx,
|
|
pred_idx] & readout_mask[patient_idx, pred_idx]
|
|
if not np.any(pred_ok):
|
|
return {
|
|
"patient": np.empty(0, dtype=np.int32),
|
|
"target_event": np.empty(0, dtype=np.int32),
|
|
"pred_idx": np.empty(0, dtype=np.int32),
|
|
"age_bin": np.empty(0, dtype=np.int16),
|
|
"target_time": np.empty(0, dtype=np.float32),
|
|
"sort_order": np.empty(0, dtype=np.int64),
|
|
"sorted_target_event": np.empty(0, dtype=np.int32),
|
|
"raw_patient": raw_patient_idx,
|
|
"raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
|
|
"raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
|
|
"p_sex": p_sex,
|
|
"age_groups": age_groups.astype(np.float32, copy=False),
|
|
"n_patients": np.int32(time_seq.shape[0]),
|
|
}
|
|
|
|
patient_idx = patient_idx[pred_ok].astype(np.int32, copy=False)
|
|
target_idx = target_idx[pred_ok]
|
|
pred_idx = pred_idx[pred_ok].astype(np.int32, copy=False)
|
|
|
|
pred_age = time_seq[patient_idx, pred_idx].astype(np.float32, copy=False)
|
|
age_bin = np.floor((pred_age - age_start) / age_step).astype(np.int16)
|
|
age_ok = (age_bin >= 0) & (age_bin < n_age)
|
|
if not np.any(age_ok):
|
|
return {
|
|
"patient": np.empty(0, dtype=np.int32),
|
|
"target_event": np.empty(0, dtype=np.int32),
|
|
"pred_idx": np.empty(0, dtype=np.int32),
|
|
"age_bin": np.empty(0, dtype=np.int16),
|
|
"target_time": np.empty(0, dtype=np.float32),
|
|
"sort_order": np.empty(0, dtype=np.int64),
|
|
"sorted_target_event": np.empty(0, dtype=np.int32),
|
|
"raw_patient": raw_patient_idx,
|
|
"raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
|
|
"raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
|
|
"p_sex": p_sex,
|
|
"age_groups": age_groups.astype(np.float32, copy=False),
|
|
"n_patients": np.int32(time_seq.shape[0]),
|
|
}
|
|
|
|
patient_idx = patient_idx[age_ok]
|
|
target_idx = target_idx[age_ok]
|
|
pred_idx = pred_idx[age_ok]
|
|
age_bin = age_bin[age_ok]
|
|
target_event = target_event_seq[patient_idx,
|
|
target_idx].astype(np.int32, copy=False)
|
|
target_time = target_time_seq[patient_idx,
|
|
target_idx].astype(np.float32, copy=False)
|
|
|
|
sort_order = np.argsort(target_event, kind="mergesort")
|
|
sorted_target_event = target_event[sort_order]
|
|
|
|
return {
|
|
"patient": patient_idx.astype(np.int32, copy=False),
|
|
"target_event": target_event,
|
|
"pred_idx": pred_idx.astype(np.int32, copy=False),
|
|
"age_bin": age_bin.astype(np.int16, copy=False),
|
|
"target_time": target_time,
|
|
"sort_order": sort_order.astype(np.int64, copy=False),
|
|
"sorted_target_event": sorted_target_event.astype(np.int32, copy=False),
|
|
"raw_patient": raw_patient_idx,
|
|
"raw_sort_order": raw_sort_order.astype(np.int64, copy=False),
|
|
"raw_sorted_target_event": raw_sorted_target_event.astype(np.int32, copy=False),
|
|
"p_sex": p_sex,
|
|
"age_groups": age_groups.astype(np.float32, copy=False),
|
|
"n_patients": np.int32(time_seq.shape[0]),
|
|
}
|
|
|
|
|
|
def _init_auc_worker_flat(
|
|
patient: np.ndarray,
|
|
target_event: np.ndarray,
|
|
pred_idx: np.ndarray,
|
|
age_bin: np.ndarray,
|
|
target_time: np.ndarray,
|
|
sort_order: np.ndarray,
|
|
sorted_target_event: np.ndarray,
|
|
raw_patient: np.ndarray,
|
|
raw_sort_order: np.ndarray,
|
|
raw_sorted_target_event: np.ndarray,
|
|
p_sex: np.ndarray,
|
|
age_groups: np.ndarray,
|
|
n_patients: int,
|
|
):
|
|
# Prevent BLAS/OpenMP oversubscription when many worker processes are active.
|
|
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
|
os.environ.setdefault("MKL_NUM_THREADS", "1")
|
|
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
|
|
os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
|
|
|
|
_WORKER.clear()
|
|
_WORKER.update({
|
|
"patient": patient,
|
|
"target_event": target_event,
|
|
"pred_idx": pred_idx,
|
|
"age_bin": age_bin,
|
|
"target_time": target_time,
|
|
"sort_order": sort_order,
|
|
"sorted_target_event": sorted_target_event,
|
|
"raw_patient": raw_patient,
|
|
"raw_sort_order": raw_sort_order,
|
|
"raw_sorted_target_event": raw_sorted_target_event,
|
|
"p_sex": p_sex,
|
|
"age_groups": age_groups,
|
|
"n_patients": int(n_patients),
|
|
})
|
|
|
|
|
|
def _case_indices_for_token(token: int) -> np.ndarray:
|
|
sorted_target_event = _WORKER["sorted_target_event"]
|
|
sort_order = _WORKER["sort_order"]
|
|
left = np.searchsorted(sorted_target_event, int(token), side="left")
|
|
right = np.searchsorted(sorted_target_event, int(token), side="right")
|
|
if right <= left:
|
|
return np.empty(0, dtype=np.int64)
|
|
return sort_order[left:right]
|
|
|
|
|
|
def _raw_case_patients_for_token(token: int) -> np.ndarray:
|
|
raw_sorted_target_event = _WORKER["raw_sorted_target_event"]
|
|
raw_sort_order = _WORKER["raw_sort_order"]
|
|
raw_patient = _WORKER["raw_patient"]
|
|
left = np.searchsorted(raw_sorted_target_event, int(token), side="left")
|
|
right = np.searchsorted(raw_sorted_target_event, int(token), side="right")
|
|
if right <= left:
|
|
return np.empty(0, dtype=np.int32)
|
|
return np.unique(raw_patient[raw_sort_order[left:right]])
|
|
|
|
|
|
def _calibration_auc_one_disease_flat(task: Tuple[int, int]) -> List[Dict[str, Any]]:
|
|
j, token = task
|
|
patient = _WORKER["patient"]
|
|
pred_idx = _WORKER["pred_idx"]
|
|
age_bin = _WORKER["age_bin"]
|
|
target_time = _WORKER["target_time"]
|
|
p_sex = _WORKER["p_sex"]
|
|
age_groups = _WORKER["age_groups"]
|
|
n_patients = _WORKER["n_patients"]
|
|
|
|
case_idx = _case_indices_for_token(int(token))
|
|
if case_idx.size < 2:
|
|
return []
|
|
|
|
case_patients = _raw_case_patients_for_token(int(token))
|
|
if case_patients.size == 0:
|
|
return []
|
|
|
|
patient_has_case = np.zeros(n_patients, dtype=bool)
|
|
patient_has_case[case_patients] = True
|
|
|
|
# Controls follow the supplied Delphi logic: any eligible target occurrence from
|
|
# a patient who never has this disease token in the evaluated target table.
|
|
control_idx = np.flatnonzero(~patient_has_case[patient])
|
|
if control_idx.size == 0:
|
|
return []
|
|
|
|
out: List[Dict[str, Any]] = []
|
|
for b, aa in enumerate(age_groups):
|
|
case_b = case_idx[age_bin[case_idx] == b]
|
|
control_b = control_idx[age_bin[control_idx] == b]
|
|
if case_b.size == 0 or control_b.size == 0:
|
|
continue
|
|
|
|
# Match previous deterministic one-occurrence-per-patient behavior within
|
|
# each age bracket, separately for cases and controls. This avoids letting
|
|
# high-utilization patients dominate the AUC.
|
|
_, case_first = np.unique(patient[case_b], return_index=True)
|
|
_, control_first = np.unique(patient[control_b], return_index=True)
|
|
case_keep = case_b[case_first]
|
|
control_keep = control_b[control_first]
|
|
if case_keep.size == 0 or control_keep.size == 0:
|
|
continue
|
|
|
|
# Delphi2M-aligned AUC score: use disease-specific eta/logit only.
|
|
# Prediction offset filters eligible prediction tokens but does not enter the score.
|
|
case_scores = p_sex[patient[case_keep], pred_idx[case_keep], j].astype(
|
|
np.float64, copy=False)
|
|
control_scores = p_sex[patient[control_keep], pred_idx[control_keep], j].astype(
|
|
np.float64, copy=False)
|
|
if case_scores.size == 0 or control_scores.size == 0:
|
|
continue
|
|
|
|
auc_value, auc_var = get_auc_delong_var(control_scores, case_scores)
|
|
|
|
out.append({
|
|
"token": int(token),
|
|
"auc": float(auc_value),
|
|
"auc_delong": float(auc_value),
|
|
"auc_variance_delong": float(auc_var),
|
|
"age": float(aa),
|
|
"age_right": float(aa + (age_groups[1] - age_groups[0])),
|
|
"n_healthy": int(control_scores.size),
|
|
"n_diseased": int(case_scores.size),
|
|
"mean_target_time": float(np.mean(target_time[case_keep])) if case_keep.size else np.nan,
|
|
})
|
|
return out
|
|
|
|
|
|
def _calibration_auc_task_block(tasks: Sequence[Tuple[int, int]]) -> List[Dict[str, Any]]:
|
|
rows: List[Dict[str, Any]] = []
|
|
for task in tasks:
|
|
rows.extend(_calibration_auc_one_disease_flat(task))
|
|
return rows
|
|
|
|
|
|
def _split_tasks_for_workers(
|
|
tasks: Sequence[Tuple[int, int]],
|
|
effective_workers: int,
|
|
task_chunk_size: int,
|
|
) -> List[List[Tuple[int, int]]]:
|
|
if not tasks:
|
|
return []
|
|
if task_chunk_size <= 0:
|
|
# Enough chunks to keep workers busy without creating one Future per disease.
|
|
task_chunk_size = max(1, math.ceil(
|
|
len(tasks) / max(1, effective_workers * 4)))
|
|
return [list(tasks[i:i + task_chunk_size]) for i in range(0, len(tasks), task_chunk_size)]
|
|
|
|
|
|
def compute_auc_chunk_parallel(
|
|
p_chunk: np.ndarray,
|
|
arrays: Dict[str, np.ndarray],
|
|
disease_ids: Sequence[int],
|
|
sex_value: int,
|
|
sex_name: str,
|
|
age_groups: np.ndarray,
|
|
offset: float,
|
|
valid_target_min_id: int,
|
|
num_workers: int,
|
|
auc_task_chunk_size: int = 0,
|
|
) -> List[Dict[str, Any]]:
|
|
sex_mask = arrays["sex"] == sex_value
|
|
if not np.any(sex_mask):
|
|
return []
|
|
|
|
flat = _build_flat_eval_table(
|
|
p_sex=p_chunk[sex_mask],
|
|
time_seq=arrays["time_seq"][sex_mask],
|
|
target_time_seq=arrays["target_time_seq"][sex_mask],
|
|
target_event_seq=arrays["target_event_seq"][sex_mask],
|
|
padding_mask=arrays["padding_mask"][sex_mask],
|
|
readout_mask=arrays["readout_mask"][sex_mask],
|
|
offset=offset,
|
|
valid_target_min_id=valid_target_min_id,
|
|
age_groups=age_groups,
|
|
)
|
|
if flat["patient"].size == 0:
|
|
return []
|
|
|
|
# Skip diseases with no cases in this sex before sending tasks to workers.
|
|
sorted_events = flat["sorted_target_event"]
|
|
tasks = []
|
|
for j, token in enumerate(disease_ids):
|
|
left = np.searchsorted(sorted_events, int(token), side="left")
|
|
right = np.searchsorted(sorted_events, int(token), side="right")
|
|
if right - left >= 2:
|
|
tasks.append((j, int(token)))
|
|
if not tasks:
|
|
return []
|
|
|
|
effective_workers = max(1, min(int(num_workers), len(tasks)))
|
|
if effective_workers <= 1:
|
|
_init_auc_worker_flat(
|
|
flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"],
|
|
flat["target_time"], flat["sort_order"], flat["sorted_target_event"],
|
|
flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"],
|
|
flat["p_sex"], flat["age_groups"], int(flat["n_patients"]),
|
|
)
|
|
nested = [_calibration_auc_one_disease_flat(t) for t in tqdm(
|
|
tasks, desc=f"AUC {sex_name}", leave=False, dynamic_ncols=True)]
|
|
else:
|
|
ctx = mp.get_context("fork") if hasattr(
|
|
os, "fork") else mp.get_context()
|
|
task_blocks = _split_tasks_for_workers(
|
|
tasks, effective_workers, int(auc_task_chunk_size))
|
|
with ProcessPoolExecutor(
|
|
max_workers=effective_workers,
|
|
mp_context=ctx,
|
|
initializer=_init_auc_worker_flat,
|
|
initargs=(
|
|
flat["patient"], flat["target_event"], flat["pred_idx"], flat["age_bin"],
|
|
flat["target_time"], flat["sort_order"], flat["sorted_target_event"],
|
|
flat["raw_patient"], flat["raw_sort_order"], flat["raw_sorted_target_event"],
|
|
flat["p_sex"], flat["age_groups"], int(flat["n_patients"]),
|
|
),
|
|
) as ex:
|
|
nested = list(tqdm(
|
|
ex.map(_calibration_auc_task_block, task_blocks),
|
|
total=len(task_blocks),
|
|
desc=f"AUC {sex_name}",
|
|
leave=False,
|
|
dynamic_ncols=True,
|
|
))
|
|
|
|
out: List[Dict[str, Any]] = []
|
|
for rows in nested:
|
|
for r in rows:
|
|
r["sex"] = sex_name
|
|
r["offset"] = float(offset)
|
|
out.append(r)
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def evaluate_auc_pipeline(
|
|
model: DeepHealth,
|
|
loader: DataLoader,
|
|
dataset: HealthDataset,
|
|
output_path: Optional[str],
|
|
diseases_of_interest: Optional[Sequence[int]],
|
|
filter_min_total: int,
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
include_death: bool,
|
|
exclude_death: bool,
|
|
disease_chunk_size: int,
|
|
age_groups: np.ndarray,
|
|
offsets: Sequence[float],
|
|
device: torch.device,
|
|
model_target_mode: str,
|
|
readout_name: str,
|
|
readout_reduce: str,
|
|
num_workers_auc: int,
|
|
use_amp: bool,
|
|
auc_task_chunk_size: int = 0,
|
|
hidden_cache_dtype: str = "float16",
|
|
logit_batch_size: int = 256,
|
|
) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
|
model.eval().to(device)
|
|
|
|
disease_ids = select_disease_tokens(
|
|
dataset=dataset,
|
|
requested_tokens=diseases_of_interest,
|
|
filter_min_total=filter_min_total,
|
|
first_occurrence_by_token=first_occurrence_by_token,
|
|
)
|
|
disease_ids = [int(k) for k in disease_ids if 0 <=
|
|
int(k) < dataset.vocab_size]
|
|
|
|
death_token_ids = _get_death_token_ids(dataset)
|
|
if (not bool(include_death)) or bool(exclude_death):
|
|
before = len(disease_ids)
|
|
death_set = set(int(x) for x in death_token_ids)
|
|
disease_ids = [int(x) for x in disease_ids if int(x) not in death_set]
|
|
print(
|
|
"[INFO] Death exclusion applied on final disease_ids: "
|
|
f"include_death={bool(include_death)}, exclude_death={bool(exclude_death)}, "
|
|
f"before={before}, after={len(disease_ids)}.")
|
|
|
|
if not disease_ids:
|
|
raise ValueError("No diseases selected for evaluation.")
|
|
|
|
if disease_chunk_size is None or int(disease_chunk_size) <= 0:
|
|
disease_chunk_size = len(disease_ids)
|
|
disease_chunk_size = max(1, int(disease_chunk_size))
|
|
|
|
num_chunks = math.ceil(len(disease_ids) / disease_chunk_size)
|
|
chunks = np.array_split(np.asarray(
|
|
disease_ids, dtype=np.int64), num_chunks)
|
|
print(
|
|
f"Evaluating {len(disease_ids)} disease tokens in {len(chunks)} chunk(s).")
|
|
print("Using Delphi2M-aligned rate/logit score for AUC.")
|
|
print("AUC score = disease-specific eta at the latest eligible prediction token.")
|
|
print("Prediction offset controls eligibility only and does not enter the score.")
|
|
print(f"Evaluating prediction offsets: {', '.join(f'{x:g}' for x in offsets)} years.")
|
|
|
|
# In current dataset sex is normalized to 0/1. UKB convention after normalization: 0=female, 1=male.
|
|
sex_items = [("female", 0), ("male", 1)]
|
|
all_rows: List[Dict[str, Any]] = []
|
|
|
|
valid_target_min_id = CHECKUP_IDX if NO_EVENT_IDX >= dataset.vocab_size else CHECKUP_IDX
|
|
# If NO_EVENT exists and should not be a disease/control target, require target > NO_EVENT_IDX.
|
|
if NO_EVENT_IDX in dataset.label_id_to_code and dataset.label_id_to_code.get(NO_EVENT_IDX) == "<NO_EVENT>":
|
|
valid_target_min_id = NO_EVENT_IDX
|
|
|
|
hidden_all, arrays = infer_readout_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 readout hidden: shape={hidden_all.shape}, dtype={hidden_all.dtype}")
|
|
|
|
for chunk_idx, chunk in enumerate(tqdm(chunks, desc="Processing disease chunks", dynamic_ncols=True)):
|
|
p_chunk = compute_logits_for_disease_chunk(
|
|
model=model,
|
|
hidden_all=hidden_all,
|
|
disease_ids=chunk.tolist(),
|
|
device=device,
|
|
logit_batch_size=logit_batch_size,
|
|
use_amp=use_amp,
|
|
)
|
|
for offset in offsets:
|
|
for sex_name, sex_value in sex_items:
|
|
rows = compute_auc_chunk_parallel(
|
|
p_chunk=p_chunk,
|
|
arrays=arrays,
|
|
disease_ids=chunk.tolist(),
|
|
sex_value=sex_value,
|
|
sex_name=sex_name,
|
|
age_groups=age_groups,
|
|
offset=float(offset),
|
|
valid_target_min_id=valid_target_min_id,
|
|
num_workers=num_workers_auc,
|
|
auc_task_chunk_size=auc_task_chunk_size,
|
|
)
|
|
for r in rows:
|
|
r["disease_chunk_idx"] = int(chunk_idx)
|
|
all_rows.extend(rows)
|
|
del p_chunk
|
|
|
|
del hidden_all, arrays
|
|
df_auc_unpooled = pd.DataFrame(all_rows)
|
|
if df_auc_unpooled.empty:
|
|
raise RuntimeError(
|
|
"No AUC rows were produced. Check offset, age_groups, eval split, and disease ids.")
|
|
|
|
# Keep outputs self-contained with only evaluation fields and token->code mapping.
|
|
df_auc_unpooled["label_code"] = df_auc_unpooled["token"].map(
|
|
dataset.label_id_to_code)
|
|
|
|
print("Using DeLong method to calculate AUC confidence intervals.")
|
|
grouped = df_auc_unpooled.groupby(
|
|
["token", "label_code", "offset"], dropna=False, as_index=False)
|
|
df_auc = grouped.agg(
|
|
auc=("auc_delong", "mean"),
|
|
n_strata=("auc_delong", "size"),
|
|
n_diseased=("n_diseased", "sum"),
|
|
n_healthy=("n_healthy", "sum"),
|
|
auc_variance_sum=("auc_variance_delong", "sum"),
|
|
)
|
|
df_auc["auc_variance_delong"] = (
|
|
df_auc["auc_variance_sum"]
|
|
/ (df_auc["n_strata"].clip(lower=1).astype(np.float64) ** 2)
|
|
)
|
|
df_auc = df_auc.drop(columns=["auc_variance_sum"])
|
|
|
|
if output_path is not None:
|
|
out_dir = Path(output_path)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
df_auc.to_csv(out_dir / "df_both.csv", index=False)
|
|
df_auc_unpooled.to_csv(
|
|
out_dir / "df_auc_unpooled.csv", index=False)
|
|
|
|
return df_auc_unpooled, df_auc
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def parse_int_list(s: Any) -> Optional[List[int]]:
|
|
if s is None:
|
|
return None
|
|
if isinstance(s, (list, tuple, np.ndarray)):
|
|
return [int(x) for x in s]
|
|
text = str(s).strip()
|
|
if text == "":
|
|
return None
|
|
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(s: Any) -> Optional[List[float]]:
|
|
if s is None:
|
|
return None
|
|
if isinstance(s, (list, tuple, np.ndarray)):
|
|
return [float(x) for x in s]
|
|
text = str(s).strip()
|
|
if text == "":
|
|
return None
|
|
if text.startswith("["):
|
|
try:
|
|
values = json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"Invalid float list: {text!r}") from exc
|
|
if not isinstance(values, list):
|
|
raise ValueError(f"Expected a JSON list, got {type(values).__name__}")
|
|
return [float(x) for x in values]
|
|
return [float(x.strip()) for x in text.split(",") if x.strip()]
|
|
|
|
|
|
def make_auc_offsets(args: argparse.Namespace, cfg: Dict[str, Any]) -> List[float]:
|
|
explicit_offsets = parse_float_list(cfg_get(args, cfg, "offsets", None))
|
|
if explicit_offsets is not None:
|
|
base_offsets = explicit_offsets
|
|
else:
|
|
next_token_offset = float(cfg_get(args, cfg, "offset", 0.1))
|
|
base_offsets = [next_token_offset, 1.0, 5.0, 10.0]
|
|
|
|
offsets: List[float] = []
|
|
seen = set()
|
|
for value in base_offsets:
|
|
value = float(value)
|
|
key = round(value, 10)
|
|
if key not in seen:
|
|
offsets.append(value)
|
|
seen.add(key)
|
|
if not offsets:
|
|
raise ValueError("At least one AUC offset is required.")
|
|
return offsets
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Evaluate DeepHealth AUC")
|
|
|
|
# Simplify arguments to only include run_path and output_path
|
|
parser.add_argument("--run_path", type=str, required=True,
|
|
help="Path containing train_config.json and best_model.pt")
|
|
parser.add_argument("--output_path", type=str,
|
|
default=None, help="Defaults to run_path")
|
|
parser.add_argument("--eval_split", type=str, default=None,
|
|
choices=["train", "val", "valid",
|
|
"validation", "test", "all"],
|
|
help="Evaluation split. Defaults to 'test' unless cfg contains eval_split.")
|
|
parser.add_argument("--dataset_subset_size", type=int, default=None,
|
|
help="Optional number of patients from the selected split.")
|
|
parser.add_argument("--batch_size", type=int, default=None,
|
|
help="Inference batch size; overrides train_config.json.")
|
|
parser.add_argument("--num_workers", type=int, default=None,
|
|
help="DataLoader workers; overrides train_config.json.")
|
|
parser.add_argument("--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,
|
|
help="CPU processes for AUC computation.")
|
|
parser.add_argument("--auc_task_chunk_size", type=int, default=None,
|
|
help="Diseases per submitted CPU task block. 0/None auto-tunes.")
|
|
parser.add_argument("--hidden_cache_dtype", type=str, default=None, choices=["float16", "float32"],
|
|
help="CPU dtype for cached readout hidden states. float16 saves memory and is usually enough for AUC.")
|
|
parser.add_argument("--logit_batch_size", type=int, default=None,
|
|
help="Patient batch size for projecting cached hidden states to disease logits.")
|
|
parser.add_argument("--disease_chunk_size", type=int, default=None,
|
|
help="Number of disease logits to materialize per inference pass. <=0 means one chunk (all diseases).")
|
|
parser.add_argument("--filter_min_total", type=int, default=None,
|
|
help="Minimum metadata count for disease selection; default 0.")
|
|
parser.add_argument("--offset", type=float, default=None,
|
|
help="Next-token prediction offset in years; preserved and evaluated alongside 1, 5, and 10 years by default.")
|
|
parser.add_argument("--offsets", type=str, default=None,
|
|
help="Comma-separated prediction offsets in years. Overrides the default set of offset,1,5,10.")
|
|
parser.add_argument("--age_start", type=float, default=None)
|
|
parser.add_argument("--age_stop", type=float, default=None)
|
|
parser.add_argument("--age_step", type=float, default=None)
|
|
parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None,
|
|
help="Use CUDA autocast during inference.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Extract paths from run_path
|
|
run_path = Path(args.run_path)
|
|
config_path = run_path / "train_config.json"
|
|
model_ckpt_path = run_path / "best_model.pt"
|
|
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"train_config.json not found in {run_path}")
|
|
if not model_ckpt_path.exists():
|
|
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
|
|
|
cfg = load_json_config(str(config_path))
|
|
|
|
if args.output_path is None:
|
|
args.output_path = str(run_path)
|
|
|
|
# Load configurations from train_config.json
|
|
data_prefix = cfg.get("data_prefix", "ukb")
|
|
labels_file = cfg.get("labels_file", "labels.csv")
|
|
no_event_interval_years = cfg.get("no_event_interval_years", 5.0)
|
|
include_no_event = cfg.get("include_no_event_in_uts_target", False)
|
|
|
|
target_mode = cfg.get("target_mode", "uts")
|
|
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 = cfg.get("dist_mode", "exponential")
|
|
readout_name = cfg.get(
|
|
"readout_name", "same_time_group_end" if target_mode == "uts" else "token")
|
|
readout_reduce = cfg.get("readout_reduce", "mean")
|
|
|
|
device = resolve_eval_device(args.device)
|
|
if device.type == "cuda":
|
|
torch.backends.cudnn.benchmark = True
|
|
|
|
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=no_event_interval_years,
|
|
include_no_event_in_uts_target=include_no_event,
|
|
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
|
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
|
exposure_cache_dir=cfg.get("exposure_cache_dir", None),
|
|
mask_onset_exposure=bool(cfg.get("mask_onset_exposure", False)),
|
|
)
|
|
validate_dataset_metadata(dataset, cfg)
|
|
|
|
subset, subset_indices = make_eval_subset(dataset, args, cfg)
|
|
print(f"Dataset: {len(dataset)} samples, vocab_size={dataset.vocab_size}")
|
|
|
|
loader = DataLoader(
|
|
subset,
|
|
batch_size=int(cfg_get(args, cfg, "batch_size", 128)),
|
|
shuffle=False,
|
|
collate_fn=sequence_eval_collate_fn,
|
|
num_workers=int(cfg_get(args, cfg, "num_workers", 4)),
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0,
|
|
prefetch_factor=2 if int(
|
|
cfg_get(args, cfg, "num_workers", 4)) > 0 else None,
|
|
)
|
|
|
|
print("Building/loading model...")
|
|
state_dict = load_checkpoint_state_dict(
|
|
str(model_ckpt_path), map_location="cpu")
|
|
dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
|
|
cfg = dict(cfg)
|
|
cfg["dist_mode"] = dist_mode
|
|
cfg["model_target_mode"] = model_target_mode
|
|
print(f"Resolved dist_mode for evaluation: {dist_mode}")
|
|
print(f"Model target mode for AUC: {model_target_mode}")
|
|
print(
|
|
"AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
|
|
"dist_mode affects model loading but is not converted to horizon-specific risk probability."
|
|
)
|
|
|
|
model = build_model_from_dataset(args, cfg, dataset).to(device)
|
|
load_model_state(model, str(model_ckpt_path),
|
|
device, state_dict=state_dict)
|
|
model.eval()
|
|
|
|
age_groups = np.arange(
|
|
float(cfg_get(args, cfg, "age_start", 40.0)),
|
|
float(cfg_get(args, cfg, "age_stop", 80.0)),
|
|
float(cfg_get(args, cfg, "age_step", 5.0)),
|
|
dtype=np.float32,
|
|
)
|
|
disease_spec = cfg_get(args, cfg, "diseases_of_interest", None)
|
|
if disease_spec is None:
|
|
disease_spec = cfg.get("disease_tokens", None)
|
|
diseases = parse_int_list(disease_spec)
|
|
first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(
|
|
dataset, subset_indices)
|
|
include_death = bool(cfg_get(args, cfg, "include_death", True))
|
|
exclude_death = bool(cfg_get(args, cfg, "exclude_death", False))
|
|
auc_offsets = make_auc_offsets(args, cfg)
|
|
|
|
evaluate_auc_pipeline(
|
|
model=model,
|
|
loader=loader,
|
|
dataset=dataset,
|
|
output_path=args.output_path,
|
|
diseases_of_interest=diseases,
|
|
filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)),
|
|
first_occurrence_by_token=first_occurrence_by_token,
|
|
include_death=include_death,
|
|
exclude_death=exclude_death,
|
|
disease_chunk_size=int(cfg_get(args, cfg, "disease_chunk_size", 0)),
|
|
age_groups=age_groups,
|
|
offsets=auc_offsets,
|
|
device=device,
|
|
model_target_mode=model_target_mode,
|
|
readout_name=readout_name,
|
|
readout_reduce=readout_reduce,
|
|
num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(
|
|
1, (os.cpu_count() or 2) - 1))),
|
|
use_amp=bool(cfg_get(args, cfg, "use_amp", False)),
|
|
auc_task_chunk_size=int(cfg_get(args, cfg, "auc_task_chunk_size", 0)),
|
|
hidden_cache_dtype=str(
|
|
cfg_get(args, cfg, "hidden_cache_dtype", "float16")),
|
|
logit_batch_size=int(
|
|
cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|