1518 lines
52 KiB
Python
1518 lines
52 KiB
Python
"""Evaluate calibration, Brier score, and NLL for all-future runs.
|
|
|
|
This evaluator intentionally does not support the Delphi2M/next-token branch.
|
|
|
|
Outputs written to each run directory:
|
|
|
|
* ``df_calibration_landmark_metrics.csv``
|
|
Per disease, sex, landmark age, and horizon IPCW metrics.
|
|
* ``df_calibration_metrics.csv``
|
|
Per disease, sex, and horizon aggregation across landmark ages.
|
|
* ``df_calibration_summary.csv``
|
|
Disease/death summaries by sex and horizon.
|
|
* ``df_calibration_curve.csv``
|
|
Fixed probability-bin calibration curves.
|
|
* ``df_point_process_nll.csv``
|
|
Exact all-future test-query NLL using the training likelihood.
|
|
* ``calibration_evaluation_summary.json``
|
|
Completion marker and evaluation metadata. It is written last.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import json
|
|
import math
|
|
import os
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, 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 (
|
|
DISEASE_HISTORY_MODE_TIMED,
|
|
AllFutureHealthDataset,
|
|
all_future_collate_fn,
|
|
normalize_disease_history_mode,
|
|
)
|
|
from eval_data import (
|
|
build_first_occurrence_map,
|
|
build_model_from_dataset,
|
|
cfg_get,
|
|
load_json_config,
|
|
load_sequence_eval_dataset,
|
|
resolve_eval_device,
|
|
validate_dataset_metadata,
|
|
validate_training_mode_config,
|
|
)
|
|
from evaluate_auc_v2 import (
|
|
LandmarkDataset,
|
|
_get_death_token_ids,
|
|
_score_to_probability,
|
|
collate_landmark_fn,
|
|
infer_landmark_hidden,
|
|
load_checkpoint_state_dict,
|
|
load_model_state,
|
|
parse_float_list,
|
|
parse_int_list,
|
|
project_distribution_chunk,
|
|
resolve_dist_mode_for_checkpoint,
|
|
select_disease_tokens,
|
|
)
|
|
from losses import build_loss
|
|
from model_architectures import resolve_model_architecture
|
|
from targets import CHECKUP_IDX, PAD_IDX
|
|
from train_util import load_eid_file
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent
|
|
COMPLETION_FILE = "calibration_evaluation_summary.json"
|
|
LANDMARK_METRICS_FILE = "df_calibration_landmark_metrics.csv"
|
|
TOKEN_METRICS_FILE = "df_calibration_metrics.csv"
|
|
SUMMARY_FILE = "df_calibration_summary.csv"
|
|
CALIBRATION_CURVE_FILE = "df_calibration_curve.csv"
|
|
POINT_PROCESS_NLL_FILE = "df_point_process_nll.csv"
|
|
|
|
DEFAULT_HORIZONS = [0.1, 1.0, 5.0, 10.0]
|
|
DEFAULT_PROBABILITY_BINS = [
|
|
0.0,
|
|
0.001,
|
|
0.002,
|
|
0.005,
|
|
0.01,
|
|
0.02,
|
|
0.05,
|
|
0.1,
|
|
0.2,
|
|
0.5,
|
|
1.0,
|
|
]
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Evaluate all-future landmark calibration, IPCW Brier score, "
|
|
"fixed-horizon NLL, and point-process NLL."
|
|
)
|
|
)
|
|
parser.add_argument("--run_path", required=True)
|
|
parser.add_argument("--output_path", default=None)
|
|
parser.add_argument(
|
|
"--eval_split",
|
|
default="test",
|
|
choices=["val", "valid", "validation", "test"],
|
|
)
|
|
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", default=None)
|
|
parser.add_argument(
|
|
"--use_amp",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=None,
|
|
)
|
|
parser.add_argument(
|
|
"--hidden_cache_dtype",
|
|
choices=["float16", "float32"],
|
|
default=None,
|
|
)
|
|
parser.add_argument("--logit_batch_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("--diseases_of_interest", default=None)
|
|
parser.add_argument("--labels_meta_path", 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",
|
|
default=None,
|
|
help="Comma-separated fixed horizons in years. Default: 0.1,1,5,10.",
|
|
)
|
|
parser.add_argument("--min_history_events", type=int, default=None)
|
|
parser.add_argument(
|
|
"--min_cases",
|
|
type=int,
|
|
default=None,
|
|
help="Minimum observed events in a landmark cell. Default: 2.",
|
|
)
|
|
parser.add_argument(
|
|
"--min_controls",
|
|
type=int,
|
|
default=1,
|
|
help="Minimum known event-free controls in a landmark cell.",
|
|
)
|
|
parser.add_argument(
|
|
"--probability_bins",
|
|
default=None,
|
|
help=(
|
|
"Comma-separated fixed probability-bin edges. "
|
|
"Default: 0,.001,.002,.005,.01,.02,.05,.1,.2,.5,1."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--max_ipcw_weight",
|
|
type=float,
|
|
default=0.0,
|
|
help="Optional IPCW cap; 0 means no cap.",
|
|
)
|
|
parser.add_argument(
|
|
"--exclude_death_in_window_without_disease",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=None,
|
|
help=(
|
|
"For non-death outcomes, treat death before disease as censoring. "
|
|
"Default: true."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--point_process_nll",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Also evaluate the exact all-future training likelihood.",
|
|
)
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Recompute even when the completion marker already exists.",
|
|
)
|
|
return parser
|
|
|
|
|
|
def _normalise_eval_split(value: str) -> str:
|
|
split = str(value).lower()
|
|
if split in {"valid", "validation"}:
|
|
return "val"
|
|
if split not in {"val", "test"}:
|
|
raise ValueError(f"eval_split must be val or test, got {value!r}")
|
|
return split
|
|
|
|
|
|
def _resolve_project_path(value: str | Path) -> Path:
|
|
path = Path(value)
|
|
if path.is_absolute():
|
|
return path
|
|
direct = Path.cwd() / path
|
|
if direct.exists():
|
|
return direct
|
|
return PROJECT_ROOT / path
|
|
|
|
|
|
def _configured_eid_file(cfg: Dict[str, Any], eval_split: str) -> Optional[Path]:
|
|
key = "val_eid_file" if eval_split == "val" else "test_eid_file"
|
|
raw = cfg.get(key)
|
|
if raw in {None, ""}:
|
|
return None
|
|
path = _resolve_project_path(str(raw))
|
|
if not path.is_file():
|
|
raise FileNotFoundError(
|
|
f"Configured {key} does not exist: {path}"
|
|
)
|
|
return path
|
|
|
|
|
|
def select_sequence_eval_indices(
|
|
dataset: Any,
|
|
cfg: Dict[str, Any],
|
|
eval_split: str,
|
|
subset_size: Optional[int],
|
|
) -> tuple[np.ndarray, str, Optional[str]]:
|
|
"""Select landmark patients using the same patient split as training."""
|
|
eval_split = _normalise_eval_split(eval_split)
|
|
eid_path = _configured_eid_file(cfg, eval_split)
|
|
if eid_path is not None:
|
|
selected_eids = load_eid_file(eid_path)
|
|
indices = np.asarray(
|
|
[
|
|
index
|
|
for index, sample in enumerate(dataset.samples)
|
|
if int(sample["eid"]) in selected_eids
|
|
],
|
|
dtype=np.int64,
|
|
)
|
|
method = "eid_file"
|
|
source = str(eid_path)
|
|
else:
|
|
n_patients = len(dataset)
|
|
train_ratio = float(cfg.get("train_ratio", 0.7))
|
|
val_ratio = float(cfg.get("val_ratio", 0.15))
|
|
test_ratio = float(cfg.get("test_ratio", 0.15))
|
|
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, got {total}")
|
|
order = np.random.RandomState(int(cfg.get("seed", 42))).permutation(
|
|
n_patients
|
|
)
|
|
n_train = int(n_patients * train_ratio)
|
|
n_val = int(n_patients * val_ratio)
|
|
indices = (
|
|
order[n_train:n_train + n_val]
|
|
if eval_split == "val"
|
|
else order[n_train + n_val:]
|
|
).astype(np.int64, copy=False)
|
|
method = "random_patient_ratio"
|
|
source = None
|
|
|
|
if subset_size is not None and int(subset_size) > 0:
|
|
indices = indices[: int(subset_size)]
|
|
if indices.size == 0:
|
|
raise RuntimeError("Selected landmark evaluation split is empty.")
|
|
return indices, method, source
|
|
|
|
|
|
def build_point_process_eval_subset(
|
|
cfg: Dict[str, Any],
|
|
eval_split: str,
|
|
disease_history_mode: str,
|
|
subset_size: Optional[int],
|
|
) -> tuple[AllFutureHealthDataset, Subset, str, Optional[str]]:
|
|
dataset_split = "valid" if _normalise_eval_split(eval_split) == "val" else "test"
|
|
dataset = AllFutureHealthDataset(
|
|
data_prefix=str(cfg.get("data_prefix", "ukb")),
|
|
labels_file=str(cfg.get("labels_file", "labels.csv")),
|
|
split=dataset_split,
|
|
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
|
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
|
validation_query_seed=int(
|
|
cfg.get("all_future_validation_query_seed", 42)
|
|
),
|
|
extra_info_types=parse_int_list(cfg.get("extra_info_types")),
|
|
disease_history_mode=disease_history_mode,
|
|
)
|
|
validate_dataset_metadata(dataset, cfg)
|
|
|
|
eval_split = _normalise_eval_split(eval_split)
|
|
eid_path = _configured_eid_file(cfg, eval_split)
|
|
if eid_path is not None:
|
|
selected_eids = load_eid_file(eid_path)
|
|
query_indices = [
|
|
query_index
|
|
for query_index, (patient_index, _query_time) in enumerate(
|
|
dataset.valid_queries
|
|
)
|
|
if int(dataset.patients[int(patient_index)]["eid"]) in selected_eids
|
|
]
|
|
method = "eid_file"
|
|
source = str(eid_path)
|
|
else:
|
|
patient_count = len(dataset.patients)
|
|
train_ratio = float(cfg.get("train_ratio", 0.7))
|
|
val_ratio = float(cfg.get("val_ratio", 0.15))
|
|
test_ratio = float(cfg.get("test_ratio", 0.15))
|
|
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, got {total}")
|
|
order = np.random.RandomState(int(cfg.get("seed", 42))).permutation(
|
|
patient_count
|
|
)
|
|
n_train = int(patient_count * train_ratio)
|
|
n_val = int(patient_count * val_ratio)
|
|
patient_set = set(
|
|
int(value)
|
|
for value in (
|
|
order[n_train:n_train + n_val]
|
|
if eval_split == "val"
|
|
else order[n_train + n_val:]
|
|
)
|
|
)
|
|
query_indices = [
|
|
query_index
|
|
for query_index, (patient_index, _query_time) in enumerate(
|
|
dataset.valid_queries
|
|
)
|
|
if int(patient_index) in patient_set
|
|
]
|
|
method = "random_patient_ratio"
|
|
source = None
|
|
|
|
if subset_size is not None and int(subset_size) > 0:
|
|
query_indices = query_indices[: int(subset_size)]
|
|
if not query_indices:
|
|
raise RuntimeError("Selected point-process evaluation split is empty.")
|
|
subset = Subset(dataset, np.asarray(query_indices, dtype=np.int64))
|
|
return dataset, subset, method, source
|
|
|
|
|
|
def _sigmoid(values: np.ndarray) -> np.ndarray:
|
|
values = np.asarray(values, dtype=np.float64)
|
|
result = np.empty_like(values, dtype=np.float64)
|
|
nonnegative = values >= 0
|
|
result[nonnegative] = 1.0 / (
|
|
1.0 + np.exp(-values[nonnegative])
|
|
)
|
|
exp_values = np.exp(values[~nonnegative])
|
|
result[~nonnegative] = exp_values / (1.0 + exp_values)
|
|
return result
|
|
|
|
|
|
def fit_weighted_logistic_calibration(
|
|
probabilities: np.ndarray,
|
|
outcomes: np.ndarray,
|
|
weights: np.ndarray,
|
|
eps: float = 1e-6,
|
|
max_iter: int = 100,
|
|
) -> tuple[float, float, float]:
|
|
"""Return calibration-in-the-large, free intercept, and free slope."""
|
|
probabilities = np.asarray(probabilities, dtype=np.float64)
|
|
outcomes = np.asarray(outcomes, dtype=np.float64)
|
|
weights = np.asarray(weights, dtype=np.float64)
|
|
valid = (
|
|
np.isfinite(probabilities)
|
|
& np.isfinite(outcomes)
|
|
& np.isfinite(weights)
|
|
& (weights > 0)
|
|
)
|
|
p = np.clip(probabilities[valid], eps, 1.0 - eps)
|
|
y = outcomes[valid]
|
|
w = weights[valid]
|
|
if p.size < 3 or np.unique(y).size < 2 or np.allclose(p, p[0]):
|
|
return np.nan, np.nan, np.nan
|
|
|
|
logit_p = np.log(p) - np.log1p(-p)
|
|
|
|
intercept_only = 0.0
|
|
for _ in range(max_iter):
|
|
fitted = _sigmoid(logit_p + intercept_only)
|
|
gradient = float(np.sum(w * (y - fitted)))
|
|
information = float(np.sum(w * fitted * (1.0 - fitted)))
|
|
if information <= 1e-12:
|
|
intercept_only = np.nan
|
|
break
|
|
step = gradient / information
|
|
intercept_only += step
|
|
if abs(step) < 1e-9:
|
|
break
|
|
|
|
design = np.column_stack([np.ones_like(logit_p), logit_p])
|
|
beta = np.asarray([0.0, 1.0], dtype=np.float64)
|
|
ridge = np.eye(2, dtype=np.float64) * 1e-9
|
|
for _ in range(max_iter):
|
|
fitted = _sigmoid(design @ beta)
|
|
variance = np.clip(fitted * (1.0 - fitted), 1e-12, None)
|
|
gradient = design.T @ (w * (y - fitted))
|
|
information = design.T @ ((w * variance)[:, None] * design) + ridge
|
|
try:
|
|
step = np.linalg.solve(information, gradient)
|
|
except np.linalg.LinAlgError:
|
|
return float(intercept_only), np.nan, np.nan
|
|
beta += step
|
|
if float(np.max(np.abs(step))) < 1e-9:
|
|
break
|
|
if not np.all(np.isfinite(beta)) or float(np.max(np.abs(beta))) > 1e6:
|
|
return float(intercept_only), np.nan, np.nan
|
|
return float(intercept_only), float(beta[0]), float(beta[1])
|
|
|
|
|
|
def _censoring_km(
|
|
observed_times: np.ndarray,
|
|
censor_events: np.ndarray,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
observed_times = np.asarray(observed_times, dtype=np.float64)
|
|
censor_events = np.asarray(censor_events, dtype=bool)
|
|
if observed_times.size == 0:
|
|
return np.empty(0, dtype=np.float64), np.empty(0, dtype=np.float64)
|
|
|
|
event_times = np.unique(observed_times[censor_events])
|
|
if event_times.size == 0:
|
|
return event_times, np.empty(0, dtype=np.float64)
|
|
|
|
survival_after = np.empty(event_times.size, dtype=np.float64)
|
|
survival = 1.0
|
|
for index, time_value in enumerate(event_times):
|
|
at_risk = int(np.sum(observed_times >= time_value))
|
|
censored = int(
|
|
np.sum(censor_events & (observed_times == time_value))
|
|
)
|
|
if at_risk > 0:
|
|
survival *= 1.0 - float(censored) / float(at_risk)
|
|
survival_after[index] = survival
|
|
return event_times, survival_after
|
|
|
|
|
|
def _km_value(
|
|
event_times: np.ndarray,
|
|
survival_after: np.ndarray,
|
|
query_times: np.ndarray | float,
|
|
*,
|
|
before: bool,
|
|
) -> np.ndarray:
|
|
values = np.asarray(query_times, dtype=np.float64)
|
|
if event_times.size == 0:
|
|
return np.ones_like(values, dtype=np.float64)
|
|
side = "left" if before else "right"
|
|
positions = np.searchsorted(event_times, values, side=side) - 1
|
|
result = np.ones_like(values, dtype=np.float64)
|
|
valid = positions >= 0
|
|
result[valid] = survival_after[positions[valid]]
|
|
return result
|
|
|
|
|
|
def _cap_ipcw(weights: np.ndarray, maximum: float) -> tuple[np.ndarray, int]:
|
|
weights = np.asarray(weights, dtype=np.float64)
|
|
if maximum <= 0:
|
|
return weights, 0
|
|
clipped = weights > maximum
|
|
return np.minimum(weights, maximum), int(clipped.sum())
|
|
|
|
|
|
def compute_ipcw_cell(
|
|
probabilities: np.ndarray,
|
|
event_times: np.ndarray,
|
|
censor_times: np.ndarray,
|
|
horizon: float,
|
|
min_cases: int,
|
|
min_controls: int,
|
|
max_ipcw_weight: float,
|
|
) -> Optional[tuple[Dict[str, Any], Dict[str, np.ndarray]]]:
|
|
"""Compute IPCW fixed-horizon metrics for one at-risk landmark cell."""
|
|
p = np.clip(np.asarray(probabilities, dtype=np.float64), 1e-8, 1.0 - 1e-8)
|
|
event_times = np.asarray(event_times, dtype=np.float64)
|
|
censor_times = np.asarray(censor_times, dtype=np.float64)
|
|
if not (p.shape == event_times.shape == censor_times.shape):
|
|
raise ValueError("probabilities, event_times, and censor_times must align")
|
|
n_at_risk = int(p.size)
|
|
if n_at_risk == 0:
|
|
return None
|
|
|
|
observed_event = event_times <= censor_times
|
|
event_by_horizon = observed_event & (event_times <= float(horizon))
|
|
controls = (event_times > float(horizon)) & (
|
|
censor_times >= float(horizon)
|
|
)
|
|
n_events = int(event_by_horizon.sum())
|
|
n_controls = int(controls.sum())
|
|
if n_events < int(min_cases) or n_controls < int(min_controls):
|
|
return None
|
|
|
|
observed_times = np.minimum(event_times, censor_times)
|
|
censor_events = censor_times < event_times
|
|
km_times, km_survival = _censoring_km(observed_times, censor_events)
|
|
|
|
event_g = _km_value(
|
|
km_times,
|
|
km_survival,
|
|
event_times[event_by_horizon],
|
|
before=True,
|
|
)
|
|
control_g = _km_value(
|
|
km_times,
|
|
km_survival,
|
|
np.full(n_controls, float(horizon), dtype=np.float64),
|
|
before=False,
|
|
)
|
|
event_weights, event_clipped = _cap_ipcw(
|
|
1.0 / np.clip(event_g, 1e-8, None),
|
|
max_ipcw_weight,
|
|
)
|
|
control_weights, control_clipped = _cap_ipcw(
|
|
1.0 / np.clip(control_g, 1e-8, None),
|
|
max_ipcw_weight,
|
|
)
|
|
|
|
metric_weights = np.zeros(n_at_risk, dtype=np.float64)
|
|
outcomes = np.zeros(n_at_risk, dtype=np.float64)
|
|
metric_weights[event_by_horizon] = event_weights
|
|
metric_weights[controls] = control_weights
|
|
outcomes[event_by_horizon] = 1.0
|
|
|
|
brier_contribution = metric_weights * np.square(outcomes - p)
|
|
nll_contribution = -metric_weights * (
|
|
outcomes * np.log(p) + (1.0 - outcomes) * np.log1p(-p)
|
|
)
|
|
event_weight_sum = float(event_weights.sum())
|
|
prediction_sum = float(p.sum())
|
|
observed_rate = event_weight_sum / float(n_at_risk)
|
|
predicted_mean = prediction_sum / float(n_at_risk)
|
|
|
|
known = event_by_horizon | controls
|
|
calibration_in_large, calibration_intercept, calibration_slope = (
|
|
fit_weighted_logistic_calibration(
|
|
probabilities=p[known],
|
|
outcomes=outcomes[known],
|
|
weights=metric_weights[known],
|
|
)
|
|
)
|
|
known_outcomes = outcomes[known]
|
|
known_probabilities = p[known]
|
|
|
|
row = {
|
|
"n_at_risk": n_at_risk,
|
|
"n_events": n_events,
|
|
"n_controls": n_controls,
|
|
"n_censored_before_horizon": int(n_at_risk - n_events - n_controls),
|
|
"known_fraction": float(known.mean()),
|
|
"prediction_sum": prediction_sum,
|
|
"event_weight_sum": event_weight_sum,
|
|
"brier_ipcw_sum": float(brier_contribution.sum()),
|
|
"nll_ipcw_sum": float(nll_contribution.sum()),
|
|
"predicted_mean": predicted_mean,
|
|
"observed_rate_ipcw": observed_rate,
|
|
"expected_observed_ratio": (
|
|
predicted_mean / observed_rate if observed_rate > 0 else np.nan
|
|
),
|
|
"brier_ipcw": float(brier_contribution.sum()) / float(n_at_risk),
|
|
"nll_ipcw": float(nll_contribution.sum()) / float(n_at_risk),
|
|
"brier_complete_case": float(
|
|
np.mean(np.square(known_outcomes - known_probabilities))
|
|
),
|
|
"nll_complete_case": float(
|
|
-np.mean(
|
|
known_outcomes * np.log(known_probabilities)
|
|
+ (1.0 - known_outcomes) * np.log1p(-known_probabilities)
|
|
)
|
|
),
|
|
"calibration_in_the_large": calibration_in_large,
|
|
"calibration_intercept": calibration_intercept,
|
|
"calibration_slope": calibration_slope,
|
|
"ipcw_weight_max": float(metric_weights.max()),
|
|
"ipcw_weight_mean_known": float(metric_weights[known].mean()),
|
|
"ipcw_weights_clipped": int(event_clipped + control_clipped),
|
|
}
|
|
arrays = {
|
|
"probabilities": p,
|
|
"outcomes": outcomes,
|
|
"metric_weights": metric_weights,
|
|
"brier_contribution": brier_contribution,
|
|
"nll_contribution": nll_contribution,
|
|
"known": known,
|
|
"event_by_horizon": event_by_horizon,
|
|
"controls": controls,
|
|
}
|
|
return row, arrays
|
|
|
|
|
|
def _update_curve_accumulator(
|
|
accumulator: Dict[tuple[Any, ...], Dict[str, float]],
|
|
*,
|
|
outcome: str,
|
|
sex: str,
|
|
horizon: float,
|
|
probability_bins: np.ndarray,
|
|
arrays: Dict[str, np.ndarray],
|
|
) -> None:
|
|
probabilities = arrays["probabilities"]
|
|
bin_indices = np.searchsorted(
|
|
probability_bins,
|
|
probabilities,
|
|
side="right",
|
|
) - 1
|
|
bin_indices = np.clip(bin_indices, 0, len(probability_bins) - 2)
|
|
|
|
for bin_index in np.unique(bin_indices).tolist():
|
|
mask = bin_indices == int(bin_index)
|
|
key = (outcome, sex, float(horizon), int(bin_index))
|
|
row = accumulator[key]
|
|
row["n_predictions"] += int(mask.sum())
|
|
row["prediction_sum"] += float(probabilities[mask].sum())
|
|
row["event_weight_sum"] += float(
|
|
(arrays["metric_weights"][mask] * arrays["outcomes"][mask]).sum()
|
|
)
|
|
row["known_weight_sum"] += float(
|
|
arrays["metric_weights"][mask].sum()
|
|
)
|
|
row["brier_ipcw_sum"] += float(
|
|
arrays["brier_contribution"][mask].sum()
|
|
)
|
|
row["nll_ipcw_sum"] += float(
|
|
arrays["nll_contribution"][mask].sum()
|
|
)
|
|
row["n_events"] += int(arrays["event_by_horizon"][mask].sum())
|
|
row["n_controls"] += int(arrays["controls"][mask].sum())
|
|
|
|
|
|
def _first_time_array(
|
|
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
|
|
token: int,
|
|
patient_count: int,
|
|
) -> np.ndarray:
|
|
result = np.full(patient_count, np.inf, dtype=np.float32)
|
|
pairs = first_occurrence_by_token.get(int(token))
|
|
if pairs is not None:
|
|
patient_ids, times = pairs
|
|
result[
|
|
np.asarray(patient_ids, dtype=np.int64)
|
|
] = np.asarray(times, dtype=np.float32)
|
|
return result
|
|
|
|
|
|
def evaluate_landmark_calibration(
|
|
*,
|
|
model: Any,
|
|
loader: DataLoader,
|
|
landmark_dataset: LandmarkDataset,
|
|
disease_ids: Sequence[int],
|
|
dist_mode: str,
|
|
horizons: np.ndarray,
|
|
device: torch.device,
|
|
use_amp: bool,
|
|
hidden_cache_dtype: str,
|
|
logit_batch_size: int,
|
|
disease_chunk_size: int,
|
|
min_cases: int,
|
|
min_controls: int,
|
|
max_ipcw_weight: float,
|
|
exclude_death_competing: bool,
|
|
probability_bins: np.ndarray,
|
|
) -> 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="all_future",
|
|
use_amp=use_amp,
|
|
hidden_cache_dtype=hidden_cache_dtype,
|
|
)
|
|
print(
|
|
f"Cached landmark hidden: shape={hidden_all.shape}, "
|
|
f"dtype={hidden_all.dtype}"
|
|
)
|
|
|
|
disease_ids = [int(token) for token in disease_ids]
|
|
disease_chunk_size = (
|
|
len(disease_ids)
|
|
if int(disease_chunk_size) <= 0
|
|
else int(disease_chunk_size)
|
|
)
|
|
chunks = [
|
|
disease_ids[start:start + disease_chunk_size]
|
|
for start in range(0, len(disease_ids), disease_chunk_size)
|
|
]
|
|
|
|
patient_count = len(landmark_dataset.subset_indices)
|
|
death_tokens = set(int(value) for value in landmark_dataset.death_token_ids)
|
|
death_index = int(
|
|
getattr(model, "death_idx", getattr(model, "vocab_size", 1) - 1)
|
|
)
|
|
metric_rows: List[Dict[str, Any]] = []
|
|
curve_accumulator: Dict[
|
|
tuple[Any, ...], Dict[str, float]
|
|
] = defaultdict(
|
|
lambda: {
|
|
"n_predictions": 0,
|
|
"prediction_sum": 0.0,
|
|
"event_weight_sum": 0.0,
|
|
"known_weight_sum": 0.0,
|
|
"brier_ipcw_sum": 0.0,
|
|
"nll_ipcw_sum": 0.0,
|
|
"n_events": 0,
|
|
"n_controls": 0,
|
|
}
|
|
)
|
|
|
|
for chunk_index, chunk_ids in enumerate(
|
|
tqdm(chunks, desc="Disease chunks", dynamic_ncols=True)
|
|
):
|
|
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,
|
|
)
|
|
|
|
for column_index, token in enumerate(
|
|
tqdm(
|
|
chunk_ids,
|
|
desc=f"Calibration chunk {chunk_index}",
|
|
leave=False,
|
|
dynamic_ncols=True,
|
|
)
|
|
):
|
|
token = int(token)
|
|
label_code = str(
|
|
landmark_dataset.dataset.label_id_to_code.get(token, token)
|
|
)
|
|
outcome_name = "Death" if token in death_tokens else "Disease"
|
|
first_time = _first_time_array(
|
|
landmark_dataset.first_occurrence_by_token,
|
|
token,
|
|
patient_count,
|
|
)
|
|
token_logits = logits_chunk[:, int(column_index)]
|
|
token_rho = (
|
|
None
|
|
if rho_chunk is None
|
|
else rho_chunk[:, int(column_index)]
|
|
)
|
|
|
|
for sex_value, sex_name in ((0, "Female"), (1, "Male")):
|
|
sex_rows = row_arrays["sex"] == int(sex_value)
|
|
if not np.any(sex_rows):
|
|
continue
|
|
landmark_values = np.unique(
|
|
row_arrays["landmark_age"][sex_rows]
|
|
)
|
|
for landmark_age_raw in landmark_values.tolist():
|
|
landmark_age = float(landmark_age_raw)
|
|
stratum = sex_rows & np.isclose(
|
|
row_arrays["landmark_age"],
|
|
np.float32(landmark_age),
|
|
)
|
|
row_indices = np.flatnonzero(stratum)
|
|
patient_ids = row_arrays["patient_id"][row_indices]
|
|
token_first_time = first_time[patient_ids].astype(
|
|
np.float64,
|
|
copy=False,
|
|
)
|
|
at_risk = token_first_time > landmark_age
|
|
if not np.any(at_risk):
|
|
continue
|
|
|
|
row_indices = row_indices[at_risk]
|
|
token_first_time = token_first_time[at_risk]
|
|
followup_end = row_arrays["followup_end_time"][
|
|
row_indices
|
|
].astype(np.float64, copy=False)
|
|
death_time = row_arrays["death_time"][
|
|
row_indices
|
|
].astype(np.float64, copy=False)
|
|
|
|
effective_censor = followup_end.copy()
|
|
if exclude_death_competing and token not in death_tokens:
|
|
death_before_disease = death_time < token_first_time
|
|
effective_censor = np.where(
|
|
death_before_disease,
|
|
np.minimum(effective_censor, death_time),
|
|
effective_censor,
|
|
)
|
|
|
|
event_times = token_first_time - landmark_age
|
|
censor_times = effective_censor - landmark_age
|
|
|
|
for horizon_raw in horizons.tolist():
|
|
horizon = float(horizon_raw)
|
|
probabilities = _score_to_probability(
|
|
token_logits[row_indices],
|
|
(
|
|
None
|
|
if token_rho is None
|
|
else token_rho[row_indices]
|
|
),
|
|
score_mode="risk",
|
|
horizon=horizon,
|
|
dist_mode=dist_mode,
|
|
token=token,
|
|
death_idx=death_index,
|
|
)
|
|
result = compute_ipcw_cell(
|
|
probabilities=probabilities,
|
|
event_times=event_times,
|
|
censor_times=censor_times,
|
|
horizon=horizon,
|
|
min_cases=min_cases,
|
|
min_controls=min_controls,
|
|
max_ipcw_weight=max_ipcw_weight,
|
|
)
|
|
if result is None:
|
|
continue
|
|
row, arrays = result
|
|
metric_rows.append(
|
|
{
|
|
"token": token,
|
|
"label_code": label_code,
|
|
"outcome": outcome_name,
|
|
"sex": sex_name,
|
|
"landmark_age": landmark_age,
|
|
"horizon": horizon,
|
|
**row,
|
|
}
|
|
)
|
|
_update_curve_accumulator(
|
|
curve_accumulator,
|
|
outcome=outcome_name,
|
|
sex=sex_name,
|
|
horizon=horizon,
|
|
probability_bins=probability_bins,
|
|
arrays=arrays,
|
|
)
|
|
|
|
del logits_chunk, rho_chunk
|
|
|
|
if not metric_rows:
|
|
raise RuntimeError(
|
|
"No calibration rows were produced. Check split, landmark ages, "
|
|
"horizons, min_cases, and disease selection."
|
|
)
|
|
|
|
metrics = pd.DataFrame(metric_rows)
|
|
curve_rows: List[Dict[str, Any]] = []
|
|
for (outcome, sex, horizon, bin_index), values in sorted(
|
|
curve_accumulator.items()
|
|
):
|
|
n_predictions = int(values["n_predictions"])
|
|
curve_rows.append(
|
|
{
|
|
"outcome": outcome,
|
|
"sex": sex,
|
|
"horizon": float(horizon),
|
|
"probability_bin": int(bin_index),
|
|
"probability_lower": float(probability_bins[int(bin_index)]),
|
|
"probability_upper": float(
|
|
probability_bins[int(bin_index) + 1]
|
|
),
|
|
"n_predictions": n_predictions,
|
|
"n_events": int(values["n_events"]),
|
|
"n_controls": int(values["n_controls"]),
|
|
"predicted_mean": (
|
|
float(values["prediction_sum"]) / n_predictions
|
|
),
|
|
"observed_rate_ipcw": (
|
|
float(values["event_weight_sum"]) / n_predictions
|
|
),
|
|
"brier_ipcw": (
|
|
float(values["brier_ipcw_sum"]) / n_predictions
|
|
),
|
|
"nll_ipcw": (
|
|
float(values["nll_ipcw_sum"]) / n_predictions
|
|
),
|
|
"known_weight_sum": float(values["known_weight_sum"]),
|
|
}
|
|
)
|
|
curve = pd.DataFrame(curve_rows)
|
|
curve_all = _aggregate_curve_sexes(curve)
|
|
curve = pd.concat([curve, curve_all], ignore_index=True)
|
|
return metrics, curve
|
|
|
|
|
|
def _aggregate_curve_sexes(curve: pd.DataFrame) -> pd.DataFrame:
|
|
rows: List[Dict[str, Any]] = []
|
|
group_columns = [
|
|
"outcome",
|
|
"horizon",
|
|
"probability_bin",
|
|
"probability_lower",
|
|
"probability_upper",
|
|
]
|
|
for keys, group in curve.groupby(group_columns, sort=True, dropna=False):
|
|
(
|
|
outcome,
|
|
horizon,
|
|
probability_bin,
|
|
probability_lower,
|
|
probability_upper,
|
|
) = keys
|
|
n_predictions = int(group["n_predictions"].sum())
|
|
prediction_sum = float(
|
|
(group["predicted_mean"] * group["n_predictions"]).sum()
|
|
)
|
|
event_weight_sum = float(
|
|
(group["observed_rate_ipcw"] * group["n_predictions"]).sum()
|
|
)
|
|
brier_sum = float(
|
|
(group["brier_ipcw"] * group["n_predictions"]).sum()
|
|
)
|
|
nll_sum = float(
|
|
(group["nll_ipcw"] * group["n_predictions"]).sum()
|
|
)
|
|
rows.append(
|
|
{
|
|
"outcome": outcome,
|
|
"sex": "All",
|
|
"horizon": float(horizon),
|
|
"probability_bin": int(probability_bin),
|
|
"probability_lower": float(probability_lower),
|
|
"probability_upper": float(probability_upper),
|
|
"n_predictions": n_predictions,
|
|
"n_events": int(group["n_events"].sum()),
|
|
"n_controls": int(group["n_controls"].sum()),
|
|
"predicted_mean": prediction_sum / n_predictions,
|
|
"observed_rate_ipcw": event_weight_sum / n_predictions,
|
|
"brier_ipcw": brier_sum / n_predictions,
|
|
"nll_ipcw": nll_sum / n_predictions,
|
|
"known_weight_sum": float(group["known_weight_sum"].sum()),
|
|
}
|
|
)
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def aggregate_metric_rows(
|
|
metrics: pd.DataFrame,
|
|
group_columns: Sequence[str],
|
|
) -> pd.DataFrame:
|
|
rows: List[Dict[str, Any]] = []
|
|
for keys, group in metrics.groupby(
|
|
list(group_columns),
|
|
sort=True,
|
|
dropna=False,
|
|
):
|
|
if not isinstance(keys, tuple):
|
|
keys = (keys,)
|
|
n_at_risk = int(group["n_at_risk"].sum())
|
|
prediction_sum = float(group["prediction_sum"].sum())
|
|
event_weight_sum = float(group["event_weight_sum"].sum())
|
|
predicted_mean = prediction_sum / n_at_risk
|
|
observed_rate = event_weight_sum / n_at_risk
|
|
rows.append(
|
|
{
|
|
**dict(zip(group_columns, keys)),
|
|
"n_landmark_cells": int(len(group)),
|
|
"n_at_risk": n_at_risk,
|
|
"n_events": int(group["n_events"].sum()),
|
|
"n_controls": int(group["n_controls"].sum()),
|
|
"n_censored_before_horizon": int(
|
|
group["n_censored_before_horizon"].sum()
|
|
),
|
|
"predicted_mean": predicted_mean,
|
|
"observed_rate_ipcw": observed_rate,
|
|
"expected_observed_ratio": (
|
|
predicted_mean / observed_rate
|
|
if observed_rate > 0
|
|
else np.nan
|
|
),
|
|
"brier_ipcw": float(group["brier_ipcw_sum"].sum())
|
|
/ n_at_risk,
|
|
"nll_ipcw": float(group["nll_ipcw_sum"].sum())
|
|
/ n_at_risk,
|
|
"calibration_in_the_large_median": float(
|
|
group["calibration_in_the_large"].median()
|
|
),
|
|
"calibration_intercept_median": float(
|
|
group["calibration_intercept"].median()
|
|
),
|
|
"calibration_slope_median": float(
|
|
group["calibration_slope"].median()
|
|
),
|
|
"ipcw_weight_max": float(group["ipcw_weight_max"].max()),
|
|
"ipcw_weights_clipped": int(
|
|
group["ipcw_weights_clipped"].sum()
|
|
),
|
|
}
|
|
)
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame:
|
|
by_sex = aggregate_metric_rows(
|
|
metrics,
|
|
group_columns=["outcome", "sex", "horizon"],
|
|
)
|
|
all_sexes = aggregate_metric_rows(
|
|
metrics,
|
|
group_columns=["outcome", "horizon"],
|
|
)
|
|
all_sexes.insert(1, "sex", "All")
|
|
return (
|
|
pd.concat([by_sex, all_sexes], ignore_index=True)
|
|
.sort_values(["outcome", "sex", "horizon"], kind="stable")
|
|
.reset_index(drop=True)
|
|
)
|
|
|
|
|
|
def _build_point_process_criterion(
|
|
dist_mode: str,
|
|
death_index: int,
|
|
) -> Any:
|
|
ignored = {PAD_IDX, CHECKUP_IDX}
|
|
if dist_mode == "exponential":
|
|
return build_loss("exponential", ignored_idx=ignored)
|
|
if dist_mode == "weibull":
|
|
return build_loss("weibull", ignored_idx=ignored)
|
|
if dist_mode == "mixed":
|
|
return build_loss(
|
|
"mixed",
|
|
death_idx=death_index,
|
|
ignored_idx=ignored,
|
|
)
|
|
raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
|
|
|
|
|
|
@torch.inference_mode()
|
|
def evaluate_point_process_nll(
|
|
*,
|
|
model: Any,
|
|
loader: DataLoader,
|
|
dist_mode: str,
|
|
device: torch.device,
|
|
use_amp: bool,
|
|
) -> Dict[str, Any]:
|
|
criterion = _build_point_process_criterion(
|
|
dist_mode,
|
|
death_index=int(
|
|
getattr(
|
|
model,
|
|
"death_idx",
|
|
int(getattr(model, "vocab_size", 1)) - 1,
|
|
)
|
|
),
|
|
)
|
|
model.eval().to(device)
|
|
total_nll = 0.0
|
|
query_count = 0
|
|
future_event_count = 0
|
|
exposure_sum = 0.0
|
|
amp_enabled = bool(use_amp and device.type == "cuda")
|
|
|
|
for batch in tqdm(
|
|
loader,
|
|
desc="Point-process NLL",
|
|
leave=False,
|
|
dynamic_ncols=True,
|
|
):
|
|
batch_device = {
|
|
key: (
|
|
value.to(device, non_blocking=True)
|
|
if isinstance(value, torch.Tensor)
|
|
else value
|
|
)
|
|
for key, value in batch.items()
|
|
}
|
|
amp_context = (
|
|
torch.autocast(device_type=device.type, dtype=torch.float16)
|
|
if amp_enabled
|
|
else contextlib.nullcontext()
|
|
)
|
|
with amp_context:
|
|
hidden = model(
|
|
event_seq=batch_device["event_seq"],
|
|
time_seq=batch_device["time_seq"],
|
|
sex=batch_device["sex"],
|
|
padding_mask=batch_device["padding_mask"],
|
|
t_query=batch_device["t_query"],
|
|
other_type=batch_device["other_type"],
|
|
other_value=batch_device["other_value"],
|
|
other_value_kind=batch_device["other_value_kind"],
|
|
other_time=batch_device["other_time"],
|
|
)
|
|
logits = model.calc_risk(hidden)
|
|
if dist_mode == "exponential":
|
|
loss = criterion(
|
|
logits=logits,
|
|
targets=batch_device["future_targets"],
|
|
exposure=batch_device["exposure"],
|
|
)
|
|
elif dist_mode == "weibull":
|
|
loss = criterion(
|
|
logits=logits,
|
|
weibull_rho=model.calc_weibull_rho(hidden),
|
|
targets=batch_device["future_targets"],
|
|
dt=batch_device["future_dt"],
|
|
exposure=batch_device["exposure"],
|
|
)
|
|
else:
|
|
loss = criterion(
|
|
logits=logits,
|
|
death_rho=model.calc_death_rho(hidden),
|
|
targets=batch_device["future_targets"],
|
|
dt=batch_device["future_dt"],
|
|
exposure=batch_device["exposure"],
|
|
)
|
|
|
|
if not torch.isfinite(loss):
|
|
raise RuntimeError("Non-finite point-process NLL encountered.")
|
|
batch_size = int(batch_device["event_seq"].shape[0])
|
|
total_nll += float(loss.detach().cpu()) * batch_size
|
|
query_count += batch_size
|
|
valid_targets = batch["future_targets"] > PAD_IDX
|
|
valid_targets &= batch["future_targets"] != CHECKUP_IDX
|
|
future_event_count += int(valid_targets.sum().item())
|
|
exposure_sum += float(batch["exposure"].sum().item())
|
|
|
|
if query_count == 0:
|
|
raise RuntimeError("Point-process NLL loader produced no queries.")
|
|
return {
|
|
"query_count": query_count,
|
|
"future_event_count": future_event_count,
|
|
"total_point_process_nll": total_nll,
|
|
"mean_point_process_nll_per_query": total_nll / query_count,
|
|
"mean_point_process_nll_per_future_event": (
|
|
total_nll / future_event_count
|
|
if future_event_count > 0
|
|
else np.nan
|
|
),
|
|
"mean_exposure_years": exposure_sum / query_count,
|
|
}
|
|
|
|
|
|
def _atomic_csv(frame: pd.DataFrame, path: Path) -> None:
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
frame.to_csv(temporary, index=False)
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def _atomic_json(payload: Dict[str, Any], path: Path) -> None:
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
temporary.write_text(
|
|
json.dumps(payload, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def main() -> None:
|
|
parser = build_parser()
|
|
args = parser.parse_args()
|
|
|
|
run_path = Path(args.run_path).resolve()
|
|
config_path = run_path / "train_config.json"
|
|
checkpoint_path = run_path / "best_model.pt"
|
|
if not config_path.is_file():
|
|
raise FileNotFoundError(config_path)
|
|
if not checkpoint_path.is_file():
|
|
raise FileNotFoundError(checkpoint_path)
|
|
|
|
cfg = load_json_config(config_path)
|
|
validate_training_mode_config(cfg)
|
|
model_target_mode = str(
|
|
cfg.get("model_target_mode", "next_token")
|
|
).lower()
|
|
if model_target_mode != "all_future":
|
|
raise ValueError(
|
|
"evaluate_calibration.py supports all_future runs only; "
|
|
f"got model_target_mode={model_target_mode!r}"
|
|
)
|
|
|
|
output_path = Path(args.output_path or run_path).resolve()
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
completion_path = output_path / COMPLETION_FILE
|
|
if completion_path.is_file() and completion_path.stat().st_size > 0:
|
|
if not args.force:
|
|
print(f"[SKIP] Existing completion marker: {completion_path}")
|
|
return
|
|
|
|
eval_split = _normalise_eval_split(args.eval_split)
|
|
subset_size = cfg_get(args, cfg, "dataset_subset_size", None)
|
|
disease_history_mode = normalize_disease_history_mode(
|
|
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
|
)
|
|
horizons = np.asarray(
|
|
parse_float_list(
|
|
cfg_get(args, cfg, "horizons", None)
|
|
)
|
|
or DEFAULT_HORIZONS,
|
|
dtype=np.float32,
|
|
)
|
|
if horizons.size == 0 or np.any(horizons <= 0):
|
|
raise ValueError("horizons must contain positive values")
|
|
probability_bins = np.asarray(
|
|
parse_float_list(args.probability_bins)
|
|
or DEFAULT_PROBABILITY_BINS,
|
|
dtype=np.float64,
|
|
)
|
|
if (
|
|
probability_bins.size < 2
|
|
or probability_bins[0] != 0.0
|
|
or probability_bins[-1] != 1.0
|
|
or np.any(np.diff(probability_bins) <= 0)
|
|
):
|
|
raise ValueError(
|
|
"probability_bins must be strictly increasing from 0 to 1"
|
|
)
|
|
|
|
print("Loading all-future sequence evaluation dataset...")
|
|
dataset = load_sequence_eval_dataset(
|
|
model_target_mode="all_future",
|
|
data_prefix=str(cfg.get("data_prefix", "ukb")),
|
|
labels_file=str(cfg.get("labels_file", "labels.csv")),
|
|
no_event_interval_years=float(
|
|
cfg.get("no_event_interval_years", 5.0)
|
|
),
|
|
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")),
|
|
disease_history_mode=disease_history_mode,
|
|
)
|
|
validate_dataset_metadata(dataset, cfg)
|
|
subset_indices, split_method, split_source = select_sequence_eval_indices(
|
|
dataset,
|
|
cfg,
|
|
eval_split,
|
|
subset_size,
|
|
)
|
|
first_occurrence_by_token = build_first_occurrence_map(
|
|
dataset,
|
|
subset_indices,
|
|
)
|
|
|
|
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:
|
|
resolved_labels_meta = _resolve_project_path(str(labels_meta_path))
|
|
if resolved_labels_meta.is_file():
|
|
labels_meta = pd.read_csv(resolved_labels_meta)
|
|
|
|
requested_diseases = parse_int_list(args.diseases_of_interest)
|
|
disease_ids = select_disease_tokens(
|
|
dataset=dataset,
|
|
labels_meta=labels_meta,
|
|
requested_tokens=requested_diseases,
|
|
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.")
|
|
|
|
state_dict = load_checkpoint_state_dict(
|
|
checkpoint_path,
|
|
map_location="cpu",
|
|
)
|
|
dist_mode = resolve_dist_mode_for_checkpoint(
|
|
str(cfg.get("dist_mode", "exponential")),
|
|
state_dict,
|
|
)
|
|
cfg_model = dict(cfg)
|
|
cfg_model["dist_mode"] = dist_mode
|
|
cfg_model["model_architecture"] = resolve_model_architecture(
|
|
cfg_model,
|
|
state_dict,
|
|
)
|
|
device = resolve_eval_device(args.device)
|
|
model = build_model_from_dataset(
|
|
args,
|
|
cfg_model,
|
|
dataset,
|
|
state_dict=state_dict,
|
|
).to(device)
|
|
load_model_state(model, state_dict)
|
|
|
|
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("No landmark ages selected.")
|
|
|
|
death_token_ids = _get_death_token_ids(dataset, labels_meta)
|
|
landmark_dataset = LandmarkDataset(
|
|
dataset=dataset,
|
|
subset_indices=subset_indices,
|
|
landmark_ages=landmark_ages,
|
|
model_target_mode="all_future",
|
|
min_history_events=int(
|
|
cfg_get(args, cfg, "min_history_events", 1)
|
|
),
|
|
first_occurrence_by_token=first_occurrence_by_token,
|
|
death_token_ids=death_token_ids,
|
|
disease_history_mode=disease_history_mode,
|
|
)
|
|
|
|
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
|
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
|
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)
|
|
)
|
|
disease_chunk_size = int(
|
|
cfg_get(args, cfg, "disease_chunk_size", 64)
|
|
)
|
|
min_cases = int(cfg_get(args, cfg, "min_cases", 2))
|
|
min_controls = int(args.min_controls)
|
|
if min_cases < 1:
|
|
raise ValueError("min_cases must be >= 1")
|
|
if min_controls < 1:
|
|
raise ValueError("min_controls must be >= 1")
|
|
if float(args.max_ipcw_weight) < 0:
|
|
raise ValueError("max_ipcw_weight must be >= 0")
|
|
exclude_death_competing = bool(
|
|
cfg_get(
|
|
args,
|
|
cfg,
|
|
"exclude_death_in_window_without_disease",
|
|
True,
|
|
)
|
|
)
|
|
|
|
landmark_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,
|
|
)
|
|
|
|
print(f"Run: {run_path}")
|
|
print(f"Eval split: {eval_split} ({split_method})")
|
|
print(f"Selected patients: {len(subset_indices)}")
|
|
print(f"Landmark queries: {len(landmark_dataset)}")
|
|
print(f"Disease tokens: {len(disease_ids)}")
|
|
print(f"Dist mode: {dist_mode}")
|
|
print(f"Disease history mode: {disease_history_mode}")
|
|
print(f"Horizons: {horizons.tolist()}")
|
|
|
|
landmark_metrics, calibration_curve = evaluate_landmark_calibration(
|
|
model=model,
|
|
loader=landmark_loader,
|
|
landmark_dataset=landmark_dataset,
|
|
disease_ids=disease_ids,
|
|
dist_mode=dist_mode,
|
|
horizons=horizons,
|
|
device=device,
|
|
use_amp=use_amp,
|
|
hidden_cache_dtype=hidden_cache_dtype,
|
|
logit_batch_size=logit_batch_size,
|
|
disease_chunk_size=disease_chunk_size,
|
|
min_cases=min_cases,
|
|
min_controls=min_controls,
|
|
max_ipcw_weight=float(args.max_ipcw_weight),
|
|
exclude_death_competing=exclude_death_competing,
|
|
probability_bins=probability_bins,
|
|
)
|
|
token_metrics = aggregate_metric_rows(
|
|
landmark_metrics,
|
|
group_columns=[
|
|
"token",
|
|
"label_code",
|
|
"outcome",
|
|
"sex",
|
|
"horizon",
|
|
],
|
|
)
|
|
calibration_summary = build_calibration_summary(landmark_metrics)
|
|
|
|
point_process_row: Dict[str, Any] = {
|
|
"evaluated": False,
|
|
"eval_split": eval_split,
|
|
}
|
|
if args.point_process_nll:
|
|
(
|
|
point_dataset,
|
|
point_subset,
|
|
point_split_method,
|
|
point_split_source,
|
|
) = build_point_process_eval_subset(
|
|
cfg,
|
|
eval_split,
|
|
disease_history_mode,
|
|
subset_size,
|
|
)
|
|
point_loader = DataLoader(
|
|
point_subset,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
collate_fn=all_future_collate_fn,
|
|
num_workers=num_workers,
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=num_workers > 0,
|
|
prefetch_factor=2 if num_workers > 0 else None,
|
|
)
|
|
point_process_row = {
|
|
"evaluated": True,
|
|
"eval_split": eval_split,
|
|
"split_method": point_split_method,
|
|
"split_source": point_split_source,
|
|
**evaluate_point_process_nll(
|
|
model=model,
|
|
loader=point_loader,
|
|
dist_mode=dist_mode,
|
|
device=device,
|
|
use_amp=use_amp,
|
|
),
|
|
}
|
|
del point_dataset, point_subset, point_loader
|
|
|
|
landmark_path = output_path / LANDMARK_METRICS_FILE
|
|
token_path = output_path / TOKEN_METRICS_FILE
|
|
summary_path = output_path / SUMMARY_FILE
|
|
curve_path = output_path / CALIBRATION_CURVE_FILE
|
|
nll_path = output_path / POINT_PROCESS_NLL_FILE
|
|
|
|
_atomic_csv(landmark_metrics, landmark_path)
|
|
_atomic_csv(token_metrics, token_path)
|
|
_atomic_csv(calibration_summary, summary_path)
|
|
_atomic_csv(calibration_curve, curve_path)
|
|
_atomic_csv(pd.DataFrame([point_process_row]), nll_path)
|
|
|
|
completion = {
|
|
"status": "complete",
|
|
"run_path": str(run_path),
|
|
"output_path": str(output_path),
|
|
"model_target_mode": model_target_mode,
|
|
"model_architecture": cfg_model["model_architecture"],
|
|
"dist_mode": dist_mode,
|
|
"time_mode": str(cfg.get("time_mode", "")),
|
|
"disease_history_mode": disease_history_mode,
|
|
"extra_info_types": parse_int_list(cfg.get("extra_info_types")) or [],
|
|
"seed": int(cfg.get("seed", 42)),
|
|
"eval_split": eval_split,
|
|
"split_method": split_method,
|
|
"split_source": split_source,
|
|
"selected_patient_count": int(len(subset_indices)),
|
|
"landmark_query_count": int(len(landmark_dataset)),
|
|
"disease_token_count": int(len(disease_ids)),
|
|
"landmark_ages": [float(value) for value in landmark_ages],
|
|
"horizons": [float(value) for value in horizons],
|
|
"min_cases": min_cases,
|
|
"min_controls": min_controls,
|
|
"exclude_death_competing": exclude_death_competing,
|
|
"landmark_estimand": (
|
|
"first-onset fixed-horizon risk; for non-death outcomes, "
|
|
"death before disease is treated as censoring"
|
|
if exclude_death_competing
|
|
else "first-onset fixed-horizon risk without death censoring"
|
|
),
|
|
"landmark_nll_definition": (
|
|
"IPCW binary negative log-likelihood at each fixed horizon"
|
|
),
|
|
"point_process_nll_definition": (
|
|
"exact all-future continuous-time training likelihood "
|
|
"on deterministic evaluation queries"
|
|
),
|
|
"max_ipcw_weight": float(args.max_ipcw_weight),
|
|
"probability_bins": [
|
|
float(value) for value in probability_bins
|
|
],
|
|
"landmark_metric_rows": int(len(landmark_metrics)),
|
|
"token_metric_rows": int(len(token_metrics)),
|
|
"summary_rows": int(len(calibration_summary)),
|
|
"calibration_curve_rows": int(len(calibration_curve)),
|
|
"point_process_nll": point_process_row,
|
|
"outputs": {
|
|
"landmark_metrics": str(landmark_path),
|
|
"token_metrics": str(token_path),
|
|
"summary": str(summary_path),
|
|
"calibration_curve": str(curve_path),
|
|
"point_process_nll": str(nll_path),
|
|
},
|
|
}
|
|
_atomic_json(completion, completion_path)
|
|
|
|
print(f"Saved: {landmark_path}")
|
|
print(f"Saved: {token_path}")
|
|
print(f"Saved: {summary_path}")
|
|
print(f"Saved: {curve_path}")
|
|
print(f"Saved: {nll_path}")
|
|
print(f"Completion marker: {completion_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|