From e471de030d93b63f1fb4e5d588ea743523cebc17 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 30 Jul 2026 17:21:57 +0800 Subject: [PATCH] Add all-future calibration evaluation --- evaluate_calibration.py | 1517 ++++++++++++++++++++++++ evaluate_calibration_all_runs_linux.sh | 379 ++++++ tests/test_calibration_metrics.py | 143 +++ 3 files changed, 2039 insertions(+) create mode 100644 evaluate_calibration.py create mode 100644 evaluate_calibration_all_runs_linux.sh create mode 100644 tests/test_calibration_metrics.py diff --git a/evaluate_calibration.py b/evaluate_calibration.py new file mode 100644 index 0000000..128c622 --- /dev/null +++ b/evaluate_calibration.py @@ -0,0 +1,1517 @@ +"""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() diff --git a/evaluate_calibration_all_runs_linux.sh b/evaluate_calibration_all_runs_linux.sh new file mode 100644 index 0000000..66ca2b3 --- /dev/null +++ b/evaluate_calibration_all_runs_linux.sh @@ -0,0 +1,379 @@ +#!/usr/bin/env bash +# +# Recursively evaluate calibration for every completed all_future run. +# +# A runnable run contains: +# - train_config.json with model_target_mode="all_future" +# - best_model.pt +# +# calibration_evaluation_summary.json is written last by the evaluator and is +# used as the completion marker. Existing non-empty markers are skipped unless +# --force is supplied. next_token/Delphi2M runs are intentionally skipped. +# +# Jobs assigned to the same GPU run sequentially; different GPUs run in +# parallel. +# +# Examples: +# bash evaluate_calibration_all_runs_linux.sh --gpus 0 +# bash evaluate_calibration_all_runs_linux.sh --gpus 0,1,2,3 +# bash evaluate_calibration_all_runs_linux.sh --gpus 0,1 --dry-run +# + +set -uo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +RUNS_ROOT="$SCRIPT_DIR/runs" +LOG_ROOT="$SCRIPT_DIR/batch_logs/evaluate_calibration_all_runs" +GPU_CSV="0" +PYTHON_BIN="${PYTHON_BIN:-python}" +NUM_WORKERS=4 +BATCH_SIZE=128 +DISEASE_CHUNK_SIZE=64 +HORIZONS="" +USE_AMP=0 +FORCE=0 +DRY_RUN=0 + +COMPLETION_FILE="calibration_evaluation_summary.json" +EVALUATOR="$SCRIPT_DIR/evaluate_calibration.py" + +usage() { + cat <<'EOF' +Usage: + bash evaluate_calibration_all_runs_linux.sh [options] + +Options: + --gpus LIST Comma-separated GPU ids (default: 0). + --runs-root PATH Root directory scanned recursively + (default: ./runs). + --log-root PATH Evaluation log root. + --python PATH Python executable + (default: $PYTHON_BIN or python). + --num-workers N DataLoader workers per job (default: 4). + --batch-size N Evaluation batch size (default: 128). + --disease-chunk-size N Disease projection chunk size (default: 64). + --horizons LIST Optional comma-separated horizons in years. + --use-amp Force CUDA automatic mixed precision. + --force Recompute runs with an existing completion file. + --dry-run Discover and print pending jobs only. + -h, --help Show this help message. + +Only all_future runs are evaluated. next_token runs are skipped. +Completion marker: calibration_evaluation_summary.json +EOF +} + +while (($# > 0)); do + case "$1" in + --gpus) + [[ $# -ge 2 ]] || { + echo "ERROR: --gpus requires a value." >&2 + exit 2 + } + GPU_CSV="$2" + shift 2 + ;; + --runs-root) + [[ $# -ge 2 ]] || { + echo "ERROR: --runs-root requires a value." >&2 + exit 2 + } + RUNS_ROOT="$2" + shift 2 + ;; + --log-root) + [[ $# -ge 2 ]] || { + echo "ERROR: --log-root requires a value." >&2 + exit 2 + } + LOG_ROOT="$2" + shift 2 + ;; + --python) + [[ $# -ge 2 ]] || { + echo "ERROR: --python requires a value." >&2 + exit 2 + } + PYTHON_BIN="$2" + shift 2 + ;; + --num-workers) + [[ $# -ge 2 ]] || { + echo "ERROR: --num-workers requires a value." >&2 + exit 2 + } + NUM_WORKERS="$2" + shift 2 + ;; + --batch-size) + [[ $# -ge 2 ]] || { + echo "ERROR: --batch-size requires a value." >&2 + exit 2 + } + BATCH_SIZE="$2" + shift 2 + ;; + --disease-chunk-size) + [[ $# -ge 2 ]] || { + echo "ERROR: --disease-chunk-size requires a value." >&2 + exit 2 + } + DISEASE_CHUNK_SIZE="$2" + shift 2 + ;; + --horizons) + [[ $# -ge 2 ]] || { + echo "ERROR: --horizons requires a value." >&2 + exit 2 + } + HORIZONS="$2" + shift 2 + ;; + --use-amp) + USE_AMP=1 + shift + ;; + --force) + FORCE=1 + shift + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "ERROR: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +[[ -n "$GPU_CSV" ]] || { + echo "ERROR: --gpus must not be empty." >&2 + exit 2 +} +[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || { + echo "ERROR: --num-workers must be a non-negative integer." >&2 + exit 2 +} +[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || { + echo "ERROR: --batch-size must be a positive integer." >&2 + exit 2 +} +[[ "$DISEASE_CHUNK_SIZE" =~ ^[1-9][0-9]*$ ]] || { + echo "ERROR: --disease-chunk-size must be a positive integer." >&2 + exit 2 +} +[[ -d "$RUNS_ROOT" ]] || { + echo "ERROR: runs root does not exist: $RUNS_ROOT" >&2 + exit 2 +} +[[ -f "$EVALUATOR" ]] || { + echo "ERROR: missing evaluator: $EVALUATOR" >&2 + exit 2 +} +command -v "$PYTHON_BIN" >/dev/null 2>&1 || { + echo "ERROR: Python executable not found: $PYTHON_BIN" >&2 + exit 2 +} + +RUNS_ROOT="$(cd -- "$RUNS_ROOT" && pwd)" +if [[ "$LOG_ROOT" != /* ]]; then + LOG_ROOT="$SCRIPT_DIR/$LOG_ROOT" +fi + +IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV" +declare -A SEEN_GPUS=() +for gpu in "${GPU_IDS[@]}"; do + [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || { + echo "ERROR: invalid GPU id: $gpu" >&2 + exit 2 + } + [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || { + echo "ERROR: duplicate GPU id: $gpu" >&2 + exit 2 + } + SEEN_GPUS["$gpu"]=1 +done + +declare -a JOB_RUN_DIRS=() +declare -a JOB_LOG_FILES=() + +add_job() { + local run_dir="$1" + local relative_run + if [[ "$run_dir" == "$RUNS_ROOT" ]]; then + relative_run="_root" + else + relative_run="${run_dir#"$RUNS_ROOT"/}" + fi + JOB_RUN_DIRS+=("$run_dir") + JOB_LOG_FILES+=("$LOG_ROOT/$relative_run/evaluate_calibration.log") +} + +run_count=0 +incomplete_count=0 +next_token_count=0 +invalid_config_count=0 +existing_count=0 + +while IFS= read -r -d '' config_path; do + run_dir="${config_path%/train_config.json}" + ((run_count += 1)) + + if [[ ! -f "$run_dir/best_model.pt" ]]; then + echo "[SKIP] Incomplete run without best_model.pt: $run_dir" + ((incomplete_count += 1)) + continue + fi + + if ! target_mode="$( + "$PYTHON_BIN" -c \ + 'import json,sys; print(str(json.load(open(sys.argv[1], encoding="utf-8")).get("model_target_mode", "next_token")).lower())' \ + "$config_path" + )"; then + echo "[SKIP] Invalid train_config.json: $config_path" >&2 + ((invalid_config_count += 1)) + continue + fi + + if [[ "$target_mode" != "all_future" ]]; then + echo "[SKIP] model_target_mode=$target_mode: $run_dir" + ((next_token_count += 1)) + continue + fi + + if ((FORCE == 0)) && [[ -s "$run_dir/$COMPLETION_FILE" ]]; then + ((existing_count += 1)) + continue + fi + + add_job "$run_dir" +done < <(find "$RUNS_ROOT" -type f -name "train_config.json" -print0) + +if ((DRY_RUN == 0)); then + mkdir -p "$LOG_ROOT" +fi + +print_command() { + printf '%q ' "$@" + printf '\n' +} + +run_job() { + local job_index="$1" + local gpu="$2" + local run_dir="${JOB_RUN_DIRS[$job_index]}" + local log_file="${JOB_LOG_FILES[$job_index]}" + local -a command=( + "$PYTHON_BIN" + -u + "$EVALUATOR" + --run_path "$run_dir" + --output_path "$run_dir" + --eval_split test + --device cuda + --num_workers "$NUM_WORKERS" + --batch_size "$BATCH_SIZE" + --disease_chunk_size "$DISEASE_CHUNK_SIZE" + ) + + if [[ -n "$HORIZONS" ]]; then + command+=(--horizons "$HORIZONS") + fi + if ((USE_AMP)); then + command+=(--use_amp) + fi + if ((FORCE)); then + command+=(--force) + fi + + echo "[$(date '+%F %T')] START gpu=$gpu" + echo " run=$run_dir" + + if ((DRY_RUN)); then + printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu" + print_command "${command[@]}" + return 0 + fi + + mkdir -p "$(dirname -- "$log_file")" + if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \ + "${command[@]}" >"$log_file" 2>&1; then + if [[ -s "$run_dir/$COMPLETION_FILE" ]]; then + echo "[$(date '+%F %T')] DONE gpu=$gpu" + return 0 + fi + echo "[$(date '+%F %T')] FAIL gpu=$gpu" >&2 + echo " Missing completion marker: $run_dir/$COMPLETION_FILE" >&2 + echo " See: $log_file" >&2 + return 1 + else + local exit_code=$? + echo "[$(date '+%F %T')] FAIL gpu=$gpu exit=$exit_code" >&2 + echo " See: $log_file" >&2 + return "$exit_code" + fi +} + +worker() { + local slot="$1" + local gpu="${GPU_IDS[$slot]}" + local job_index + local failed=0 + + for (( + job_index = slot; + job_index < ${#JOB_RUN_DIRS[@]}; + job_index += ${#GPU_IDS[@]} + )); do + run_job "$job_index" "$gpu" || failed=1 + done + return "$failed" +} + +echo "Runs root: $RUNS_ROOT" +echo "GPUs: ${GPU_IDS[*]}" +echo "Runs discovered: $run_count" +echo "Incomplete runs skipped: $incomplete_count" +echo "next_token runs skipped: $next_token_count" +echo "Invalid configs skipped: $invalid_config_count" +echo "Existing calibration results skipped: $existing_count" +echo "Pending all_future evaluations: ${#JOB_RUN_DIRS[@]}" +echo "Log root: $LOG_ROOT" +echo + +if ((${#JOB_RUN_DIRS[@]} == 0)); then + echo "No pending all_future calibration evaluations." + exit 0 +fi + +declare -a WORKER_PIDS=() +for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do + worker "$slot" & + WORKER_PIDS+=("$!") +done + +overall_status=0 +for pid in "${WORKER_PIDS[@]}"; do + wait "$pid" || overall_status=1 +done + +if ((overall_status != 0)); then + echo "One or more calibration evaluations failed." >&2 + echo "Inspect logs under: $LOG_ROOT" >&2 + exit 1 +fi + +if ((DRY_RUN)); then + echo "Dry run completed successfully." +else + echo "All pending all_future calibration evaluations completed." +fi diff --git a/tests/test_calibration_metrics.py b/tests/test_calibration_metrics.py new file mode 100644 index 0000000..26fc110 --- /dev/null +++ b/tests/test_calibration_metrics.py @@ -0,0 +1,143 @@ +import math +import unittest + +import numpy as np +import pandas as pd + +from evaluate_calibration import ( + aggregate_metric_rows, + compute_ipcw_cell, + fit_weighted_logistic_calibration, +) + + +class IPCWCalibrationMetricTests(unittest.TestCase): + def test_no_censoring_matches_binary_metrics(self): + result = compute_ipcw_cell( + probabilities=np.asarray([0.2, 0.8]), + event_times=np.asarray([np.inf, 0.5]), + censor_times=np.asarray([2.0, 2.0]), + horizon=1.0, + min_cases=1, + min_controls=1, + max_ipcw_weight=0.0, + ) + + self.assertIsNotNone(result) + row, arrays = result + self.assertEqual(row["n_events"], 1) + self.assertEqual(row["n_controls"], 1) + self.assertAlmostEqual(row["brier_ipcw"], 0.04) + self.assertAlmostEqual(row["nll_ipcw"], -math.log(0.8)) + self.assertAlmostEqual(row["predicted_mean"], 0.5) + self.assertAlmostEqual(row["observed_rate_ipcw"], 0.5) + np.testing.assert_allclose(arrays["metric_weights"], [1.0, 1.0]) + + def test_censored_before_horizon_gets_zero_outcome_weight(self): + result = compute_ipcw_cell( + probabilities=np.asarray([0.8, 0.2, 0.4]), + event_times=np.asarray([0.5, np.inf, np.inf]), + censor_times=np.asarray([2.0, 2.0, 0.5]), + horizon=1.0, + min_cases=1, + min_controls=1, + max_ipcw_weight=0.0, + ) + + self.assertIsNotNone(result) + row, arrays = result + self.assertEqual(row["n_censored_before_horizon"], 1) + self.assertAlmostEqual(row["known_fraction"], 2.0 / 3.0) + np.testing.assert_allclose( + arrays["metric_weights"], + [1.0, 1.5, 0.0], + ) + self.assertAlmostEqual(row["brier_ipcw"], 0.1 / 3.0) + self.assertAlmostEqual( + row["nll_ipcw"], + -2.5 * math.log(0.8) / 3.0, + ) + self.assertAlmostEqual(row["observed_rate_ipcw"], 1.0 / 3.0) + + def test_calibration_intercept_and_slope_recover_identity(self): + probabilities = np.repeat([0.1, 0.3, 0.7, 0.9], 100) + outcomes = np.concatenate( + [ + np.r_[np.ones(10), np.zeros(90)], + np.r_[np.ones(30), np.zeros(70)], + np.r_[np.ones(70), np.zeros(30)], + np.r_[np.ones(90), np.zeros(10)], + ] + ) + weights = np.ones_like(probabilities) + + calibration_in_large, intercept, slope = ( + fit_weighted_logistic_calibration( + probabilities, + outcomes, + weights, + ) + ) + + self.assertAlmostEqual(calibration_in_large, 0.0, places=7) + self.assertAlmostEqual(intercept, 0.0, places=7) + self.assertAlmostEqual(slope, 1.0, places=7) + + def test_metric_aggregation_uses_contribution_sums(self): + metrics = pd.DataFrame( + [ + { + "outcome": "Disease", + "sex": "Female", + "horizon": 5.0, + "n_at_risk": 10, + "n_events": 2, + "n_controls": 7, + "n_censored_before_horizon": 1, + "prediction_sum": 2.0, + "event_weight_sum": 2.0, + "brier_ipcw_sum": 1.0, + "nll_ipcw_sum": 3.0, + "calibration_in_the_large": 0.1, + "calibration_intercept": 0.2, + "calibration_slope": 0.9, + "ipcw_weight_max": 1.2, + "ipcw_weights_clipped": 0, + }, + { + "outcome": "Disease", + "sex": "Female", + "horizon": 5.0, + "n_at_risk": 10, + "n_events": 3, + "n_controls": 6, + "n_censored_before_horizon": 1, + "prediction_sum": 3.0, + "event_weight_sum": 3.0, + "brier_ipcw_sum": 2.0, + "nll_ipcw_sum": 4.0, + "calibration_in_the_large": -0.1, + "calibration_intercept": -0.2, + "calibration_slope": 1.1, + "ipcw_weight_max": 1.4, + "ipcw_weights_clipped": 1, + }, + ] + ) + + aggregated = aggregate_metric_rows( + metrics, + group_columns=["outcome", "sex", "horizon"], + ).iloc[0] + + self.assertEqual(aggregated["n_at_risk"], 20) + self.assertAlmostEqual(aggregated["predicted_mean"], 0.25) + self.assertAlmostEqual(aggregated["observed_rate_ipcw"], 0.25) + self.assertAlmostEqual(aggregated["brier_ipcw"], 0.15) + self.assertAlmostEqual(aggregated["nll_ipcw"], 0.35) + self.assertAlmostEqual(aggregated["calibration_slope_median"], 1.0) + self.assertEqual(aggregated["ipcw_weights_clipped"], 1) + + +if __name__ == "__main__": + unittest.main()