From 4580058687c6a878f1bce85e96e3f074cd550589 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 30 Jul 2026 18:29:15 +0800 Subject: [PATCH] Optimize calibration evaluation --- evaluate_calibration.py | 1053 +++++++++++++++++------- evaluate_calibration_all_runs_linux.sh | 25 + tests/test_calibration_metrics.py | 321 ++++++++ 3 files changed, 1121 insertions(+), 278 deletions(-) diff --git a/evaluate_calibration.py b/evaluate_calibration.py index 128c622..d07423d 100644 --- a/evaluate_calibration.py +++ b/evaluate_calibration.py @@ -26,6 +26,7 @@ import json import math import os from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -54,7 +55,6 @@ from eval_data import ( from evaluate_auc_v2 import ( LandmarkDataset, _get_death_token_ids, - _score_to_probability, collate_landmark_fn, infer_landmark_hidden, load_checkpoint_state_dict, @@ -125,6 +125,15 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--logit_batch_size", type=int, default=None) parser.add_argument("--disease_chunk_size", type=int, default=None) + parser.add_argument( + "--num_workers_calibration", + type=int, + default=None, + help=( + "CPU threads for per-disease calibration statistics. " + "0 uses all logical CPUs." + ), + ) 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) @@ -360,79 +369,221 @@ def fit_weighted_logistic_calibration( max_iter: int = 100, ) -> tuple[float, float, float]: """Return calibration-in-the-large, free intercept, and free slope.""" + calibration_in_large, intercept, slope = ( + fit_weighted_logistic_calibration_batch( + probabilities=np.asarray(probabilities, dtype=np.float64)[None, :], + outcomes=np.asarray(outcomes, dtype=np.float64)[None, :], + weights=np.asarray(weights, dtype=np.float64)[None, :], + eps=eps, + max_iter=max_iter, + ) + ) + return ( + float(calibration_in_large[0]), + float(intercept[0]), + float(slope[0]), + ) + + +def fit_weighted_logistic_calibration_batch( + probabilities: np.ndarray, + outcomes: np.ndarray, + weights: np.ndarray, + eps: float = 1e-6, + max_iter: int = 100, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Fit all horizon-specific calibration models in one vectorized pass.""" probabilities = np.asarray(probabilities, dtype=np.float64) outcomes = np.asarray(outcomes, dtype=np.float64) weights = np.asarray(weights, dtype=np.float64) + if probabilities.ndim != 2: + raise ValueError("probabilities must have shape [horizons, samples]") + if not ( + probabilities.shape == outcomes.shape == weights.shape + ): + raise ValueError("probabilities, outcomes, and weights must align") + 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 - + p = np.where( + valid, + np.clip(probabilities, eps, 1.0 - eps), + 0.5, + ) + y = np.where(valid, outcomes, 0.0) + w = np.where(valid, weights, 0.0) 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 + valid_count = valid.sum(axis=1) + event_count = (valid & (y > 0.5)).sum(axis=1) + control_count = (valid & (y <= 0.5)).sum(axis=1) + probability_min = np.min( + np.where(valid, p, np.inf), + axis=1, + ) + probability_max = np.max( + np.where(valid, p, -np.inf), + axis=1, + ) + eligible = ( + (valid_count >= 3) + & (event_count > 0) + & (control_count > 0) + & ((probability_max - probability_min) > 1e-12) + ) + horizon_count = int(probabilities.shape[0]) - 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 + intercept_only = np.zeros(horizon_count, dtype=np.float64) + intercept_active = eligible.copy() 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: + if not np.any(intercept_active): 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]) + fitted = _sigmoid(logit_p + intercept_only[:, None]) + gradient = np.sum(w * (y - fitted), axis=1) + information = np.sum( + w * fitted * (1.0 - fitted), + axis=1, + ) + solvable = ( + intercept_active + & np.isfinite(gradient) + & np.isfinite(information) + & (information > 1e-12) + ) + failed = intercept_active & ~solvable + intercept_only[failed] = np.nan + intercept_active[failed] = False + step = np.zeros(horizon_count, dtype=np.float64) + step[solvable] = gradient[solvable] / information[solvable] + intercept_only[solvable] += step[solvable] + diverged = solvable & ( + ~np.isfinite(intercept_only) + | (np.abs(intercept_only) > 1e6) + ) + intercept_only[diverged] = np.nan + intercept_active[diverged] = False + intercept_active[solvable & (np.abs(step) < 1e-9)] = False + + beta_intercept = np.zeros(horizon_count, dtype=np.float64) + beta_slope = np.ones(horizon_count, dtype=np.float64) + beta_active = eligible.copy() + ridge = 1e-9 + for _ in range(max_iter): + if not np.any(beta_active): + break + fitted = _sigmoid( + beta_intercept[:, None] + beta_slope[:, None] * logit_p + ) + variance = np.clip(fitted * (1.0 - fitted), 1e-12, None) + residual = w * (y - fitted) + weighted_variance = w * variance + + gradient_0 = np.sum(residual, axis=1) + gradient_1 = np.sum(residual * logit_p, axis=1) + information_00 = np.sum(weighted_variance, axis=1) + ridge + information_01 = np.sum( + weighted_variance * logit_p, + axis=1, + ) + information_11 = np.sum( + weighted_variance * np.square(logit_p), + axis=1, + ) + ridge + determinant = ( + information_00 * information_11 + - np.square(information_01) + ) + solvable = ( + beta_active + & np.isfinite(gradient_0) + & np.isfinite(gradient_1) + & np.isfinite(determinant) + & (determinant > 1e-18) + ) + failed = beta_active & ~solvable + beta_intercept[failed] = np.nan + beta_slope[failed] = np.nan + beta_active[failed] = False + + step_0 = np.zeros(horizon_count, dtype=np.float64) + step_1 = np.zeros(horizon_count, dtype=np.float64) + step_0[solvable] = ( + gradient_0[solvable] * information_11[solvable] + - gradient_1[solvable] * information_01[solvable] + ) / determinant[solvable] + step_1[solvable] = ( + information_00[solvable] * gradient_1[solvable] + - information_01[solvable] * gradient_0[solvable] + ) / determinant[solvable] + beta_intercept[solvable] += step_0[solvable] + beta_slope[solvable] += step_1[solvable] + diverged = solvable & ( + ~np.isfinite(beta_intercept) + | ~np.isfinite(beta_slope) + | (np.abs(beta_intercept) > 1e6) + | (np.abs(beta_slope) > 1e6) + ) + beta_intercept[diverged] = np.nan + beta_slope[diverged] = np.nan + beta_active[diverged] = False + converged = solvable & ( + np.maximum(np.abs(step_0), np.abs(step_1)) < 1e-9 + ) + beta_active[converged] = False + + intercept_only[~eligible] = np.nan + beta_intercept[~eligible] = np.nan + beta_slope[~eligible] = np.nan + return intercept_only, beta_intercept, beta_slope def _censoring_km( observed_times: np.ndarray, censor_events: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: + """Estimate censoring survival in O(N log N) time.""" 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) + if observed_times.shape != censor_events.shape: + raise ValueError("observed_times and censor_events must align") - 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)) + order = np.argsort(observed_times, kind="stable") + sorted_times = observed_times[order] + sorted_censor_events = censor_events[order].astype( + np.int64, + copy=False, + ) + unique_times, first_indices = np.unique( + sorted_times, + return_index=True, + ) + censor_counts = np.add.reduceat( + sorted_censor_events, + first_indices, + ) + censor_mask = censor_counts > 0 + if not np.any(censor_mask): + return ( + np.empty(0, dtype=np.float64), + np.empty(0, dtype=np.float64), ) - if at_risk > 0: - survival *= 1.0 - float(censored) / float(at_risk) - survival_after[index] = survival + + event_times = unique_times[censor_mask] + at_risk = ( + int(observed_times.size) - first_indices[censor_mask] + ).astype(np.float64, copy=False) + survival_factors = ( + 1.0 + - censor_counts[censor_mask].astype(np.float64, copy=False) + / at_risk + ) + survival_after = np.cumprod(survival_factors, dtype=np.float64) return event_times, survival_after @@ -472,24 +623,67 @@ def compute_ipcw_cell( 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) + results = compute_ipcw_horizons( + probabilities=np.asarray(probabilities, dtype=np.float64)[None, :], + event_times=event_times, + censor_times=censor_times, + horizons=np.asarray([horizon], dtype=np.float64), + min_cases=min_cases, + min_controls=min_controls, + max_ipcw_weight=max_ipcw_weight, + ) + return results[0] + + +def compute_ipcw_horizons( + probabilities: np.ndarray, + event_times: np.ndarray, + censor_times: np.ndarray, + horizons: np.ndarray, + min_cases: int, + min_controls: int, + max_ipcw_weight: float, +) -> List[Optional[tuple[Dict[str, Any], Dict[str, np.ndarray]]]]: + """Compute all fixed horizons with one censoring KM and batched fits.""" + 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) + horizons = np.asarray(horizons, dtype=np.float64) + if p.ndim != 2: + raise ValueError("probabilities must have shape [horizons, samples]") + if horizons.ndim != 1 or p.shape[0] != horizons.size: + raise ValueError("probabilities and horizons must align") + if not ( + event_times.shape == censor_times.shape == (p.shape[1],) + ): + raise ValueError( + "probabilities, event_times, and censor_times must align" + ) + n_at_risk = int(p.shape[1]) if n_at_risk == 0: - return None + return [None] * int(horizons.size) observed_event = event_times <= censor_times - event_by_horizon = observed_event & (event_times <= float(horizon)) - controls = (event_times > float(horizon)) & ( - censor_times >= float(horizon) + event_by_horizon = ( + observed_event[None, :] + & (event_times[None, :] <= horizons[:, None]) ) - 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 + controls = ( + (event_times[None, :] > horizons[:, None]) + & (censor_times[None, :] >= horizons[:, None]) + ) + n_events = event_by_horizon.sum(axis=1).astype(np.int64) + n_controls = controls.sum(axis=1).astype(np.int64) + eligible = ( + (n_events >= int(min_cases)) + & (n_controls >= int(min_controls)) + ) + if not np.any(eligible): + return [None] * int(horizons.size) observed_times = np.minimum(event_times, censor_times) censor_events = censor_times < event_times @@ -498,94 +692,170 @@ def compute_ipcw_cell( event_g = _km_value( km_times, km_survival, - event_times[event_by_horizon], + event_times, before=True, ) control_g = _km_value( km_times, km_survival, - np.full(n_controls, float(horizon), dtype=np.float64), + horizons, before=False, ) - event_weights, event_clipped = _cap_ipcw( - 1.0 / np.clip(event_g, 1e-8, None), + raw_event_weights = 1.0 / np.clip(event_g, 1e-8, None) + raw_control_weights = 1.0 / np.clip(control_g, 1e-8, None) + event_weight_values, _ = _cap_ipcw( + raw_event_weights, max_ipcw_weight, ) - control_weights, control_clipped = _cap_ipcw( - 1.0 / np.clip(control_g, 1e-8, None), + control_weight_values, _ = _cap_ipcw( + raw_control_weights, max_ipcw_weight, ) + event_clipped = ( + raw_event_weights > max_ipcw_weight + if max_ipcw_weight > 0 + else np.zeros(n_at_risk, dtype=bool) + ) + control_clipped = ( + raw_control_weights > max_ipcw_weight + if max_ipcw_weight > 0 + else np.zeros(horizons.size, dtype=bool) + ) - 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 + metric_weights = np.where( + event_by_horizon, + event_weight_values[None, :], + 0.0, + ) + metric_weights = np.where( + controls, + control_weight_values[:, None], + metric_weights, + ) + outcomes = event_by_horizon.astype(np.float64) 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], + fit_weighted_logistic_calibration_batch( + probabilities=p, + outcomes=outcomes, + weights=metric_weights, ) ) - 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) + prediction_sum = p.sum(axis=1) + event_weight_sum = ( + metric_weights * outcomes + ).sum(axis=1) + brier_sum = brier_contribution.sum(axis=1) + nll_sum = nll_contribution.sum(axis=1) + known_count = known.sum(axis=1) + complete_case_brier = ( + np.where(known, np.square(outcomes - p), 0.0).sum(axis=1) + / np.maximum(known_count, 1) + ) + complete_case_nll = ( + -np.where( + known, + outcomes * np.log(p) + + (1.0 - outcomes) * np.log1p(-p), + 0.0, + ).sum(axis=1) + / np.maximum(known_count, 1) + ) + + results: List[ + Optional[tuple[Dict[str, Any], Dict[str, np.ndarray]]] + ] = [] + for horizon_index in range(int(horizons.size)): + if not bool(eligible[horizon_index]): + results.append(None) + continue + observed_rate = ( + float(event_weight_sum[horizon_index]) / float(n_at_risk) + ) + predicted_mean = ( + float(prediction_sum[horizon_index]) / float(n_at_risk) + ) + clipped_count = int( + np.sum( + event_by_horizon[horizon_index] & event_clipped ) - ), - "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 + + ( + int(n_controls[horizon_index]) + if bool(control_clipped[horizon_index]) + else 0 + ) + ) + row = { + "n_at_risk": n_at_risk, + "n_events": int(n_events[horizon_index]), + "n_controls": int(n_controls[horizon_index]), + "n_censored_before_horizon": int( + n_at_risk + - n_events[horizon_index] + - n_controls[horizon_index] + ), + "known_fraction": float( + known_count[horizon_index] / float(n_at_risk) + ), + "prediction_sum": float(prediction_sum[horizon_index]), + "event_weight_sum": float(event_weight_sum[horizon_index]), + "brier_ipcw_sum": float(brier_sum[horizon_index]), + "nll_ipcw_sum": float(nll_sum[horizon_index]), + "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_sum[horizon_index]) / float(n_at_risk) + ), + "nll_ipcw": ( + float(nll_sum[horizon_index]) / float(n_at_risk) + ), + "brier_complete_case": float( + complete_case_brier[horizon_index] + ), + "nll_complete_case": float( + complete_case_nll[horizon_index] + ), + "calibration_in_the_large": float( + calibration_in_large[horizon_index] + ), + "calibration_intercept": float( + calibration_intercept[horizon_index] + ), + "calibration_slope": float( + calibration_slope[horizon_index] + ), + "ipcw_weight_max": float( + metric_weights[horizon_index].max() + ), + "ipcw_weight_mean_known": float( + metric_weights[horizon_index].sum() + / known_count[horizon_index] + ), + "ipcw_weights_clipped": clipped_count, + } + arrays = { + "probabilities": p[horizon_index], + "outcomes": outcomes[horizon_index], + "metric_weights": metric_weights[horizon_index], + "brier_contribution": brier_contribution[horizon_index], + "nll_contribution": nll_contribution[horizon_index], + "known": known[horizon_index], + "event_by_horizon": event_by_horizon[horizon_index], + "controls": controls[horizon_index], + } + results.append((row, arrays)) + return results def _update_curve_accumulator( @@ -604,27 +874,81 @@ def _update_curve_accumulator( side="right", ) - 1 bin_indices = np.clip(bin_indices, 0, len(probability_bins) - 2) + bin_count = int(len(probability_bins) - 1) + n_predictions = np.bincount( + bin_indices, + minlength=bin_count, + ) + prediction_sum = np.bincount( + bin_indices, + weights=probabilities, + minlength=bin_count, + ) + event_weight_sum = np.bincount( + bin_indices, + weights=arrays["metric_weights"] * arrays["outcomes"], + minlength=bin_count, + ) + known_weight_sum = np.bincount( + bin_indices, + weights=arrays["metric_weights"], + minlength=bin_count, + ) + brier_sum = np.bincount( + bin_indices, + weights=arrays["brier_contribution"], + minlength=bin_count, + ) + nll_sum = np.bincount( + bin_indices, + weights=arrays["nll_contribution"], + minlength=bin_count, + ) + n_events = np.bincount( + bin_indices, + weights=arrays["event_by_horizon"].astype(np.int64), + minlength=bin_count, + ) + n_controls = np.bincount( + bin_indices, + weights=arrays["controls"].astype(np.int64), + minlength=bin_count, + ) - for bin_index in np.unique(bin_indices).tolist(): - mask = bin_indices == int(bin_index) + for bin_index in np.flatnonzero(n_predictions).tolist(): 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()) + row["n_predictions"] += int(n_predictions[bin_index]) + row["prediction_sum"] += float(prediction_sum[bin_index]) + row["event_weight_sum"] += float(event_weight_sum[bin_index]) + row["known_weight_sum"] += float(known_weight_sum[bin_index]) + row["brier_ipcw_sum"] += float(brier_sum[bin_index]) + row["nll_ipcw_sum"] += float(nll_sum[bin_index]) + row["n_events"] += int(n_events[bin_index]) + row["n_controls"] += int(n_controls[bin_index]) + + +def _empty_curve_values() -> Dict[str, float]: + return { + "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, + } + + +def _merge_curve_accumulator( + destination: Dict[tuple[Any, ...], Dict[str, float]], + source: Dict[tuple[Any, ...], Dict[str, float]], +) -> None: + for key, source_values in source.items(): + destination_values = destination[key] + for field, value in source_values.items(): + destination_values[field] += value def _first_time_array( @@ -642,6 +966,193 @@ def _first_time_array( return result +def _risk_probability_matrix( + logits: np.ndarray, + rho: Optional[np.ndarray], + horizons: np.ndarray, + dist_mode: str, + token: int, + death_idx: int, +) -> np.ndarray: + """Convert one token's logits to all horizon risks at once.""" + logits = np.asarray(logits, dtype=np.float32) + horizons = np.asarray(horizons, dtype=np.float32) + rate = ( + np.log1p(np.exp(-np.abs(logits))) + + np.maximum(logits, np.float32(0.0)) + + np.float32(1e-8) + ) + use_weibull = ( + str(dist_mode).lower() == "weibull" + or ( + str(dist_mode).lower() == "mixed" + and int(token) == int(death_idx) + ) + ) + if use_weibull: + if rho is None: + raise RuntimeError( + "Weibull risk scoring requires rho parameters." + ) + exposure = np.power( + horizons[:, None], + np.asarray(rho, dtype=np.float32)[None, :], + ) + else: + exposure = horizons[:, None] + return ( + -np.expm1(-rate[None, :] * exposure) + ).astype(np.float64, copy=False) + + +def _build_calibration_strata( + row_sex: np.ndarray, + row_landmark_age: np.ndarray, +) -> List[tuple[str, float, np.ndarray]]: + row_sex = np.asarray(row_sex) + row_landmark_age = np.asarray(row_landmark_age) + strata: List[tuple[str, float, np.ndarray]] = [] + for sex_value, sex_name in ((0, "Female"), (1, "Male")): + sex_rows = row_sex == int(sex_value) + if not np.any(sex_rows): + continue + for landmark_age_raw in np.unique( + row_landmark_age[sex_rows] + ).tolist(): + landmark_age = float(landmark_age_raw) + row_indices = np.flatnonzero( + sex_rows + & (row_landmark_age == np.float32(landmark_age)) + ) + if row_indices.size: + strata.append((sex_name, landmark_age, row_indices)) + return strata + + +def _evaluate_calibration_token( + *, + column_index: int, + token: int, + logits_chunk: np.ndarray, + rho_chunk: Optional[np.ndarray], + strata: Sequence[tuple[str, float, np.ndarray]], + row_patient_id: np.ndarray, + row_followup_end: np.ndarray, + row_death_time: np.ndarray, + first_occurrence_by_token: Dict[ + int, Tuple[np.ndarray, np.ndarray] + ], + patient_count: int, + death_tokens: set[int], + label_id_to_code: Dict[int, str], + dist_mode: str, + horizons: np.ndarray, + death_index: int, + min_cases: int, + min_controls: int, + max_ipcw_weight: float, + exclude_death_competing: bool, + probability_bins: np.ndarray, +) -> tuple[ + List[Dict[str, Any]], + Dict[tuple[Any, ...], Dict[str, float]], +]: + token = int(token) + token_logits = logits_chunk[:, int(column_index)] + token_rho = ( + None + if rho_chunk is None + else rho_chunk[:, int(column_index)] + ) + first_time = _first_time_array( + first_occurrence_by_token, + token, + patient_count, + ) + label_code = str(label_id_to_code.get(token, token)) + outcome_name = "Death" if token in death_tokens else "Disease" + metric_rows: List[Dict[str, Any]] = [] + curve_accumulator: Dict[ + tuple[Any, ...], Dict[str, float] + ] = defaultdict(_empty_curve_values) + + for sex_name, landmark_age, stratum_indices in strata: + patient_ids = row_patient_id[stratum_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 = stratum_indices[at_risk] + token_first_time = token_first_time[at_risk] + followup_end = row_followup_end[row_indices].astype( + np.float64, + copy=False, + ) + death_time = row_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, + ) + + probabilities = _risk_probability_matrix( + logits=token_logits[row_indices], + rho=( + None + if token_rho is None + else token_rho[row_indices] + ), + horizons=horizons, + dist_mode=dist_mode, + token=token, + death_idx=death_index, + ) + results = compute_ipcw_horizons( + probabilities=probabilities, + event_times=token_first_time - landmark_age, + censor_times=effective_censor - landmark_age, + horizons=horizons, + min_cases=min_cases, + min_controls=min_controls, + max_ipcw_weight=max_ipcw_weight, + ) + for horizon, result in zip(horizons.tolist(), results): + 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": float(horizon), + **row, + } + ) + _update_curve_accumulator( + curve_accumulator, + outcome=outcome_name, + sex=sex_name, + horizon=float(horizon), + probability_bins=probability_bins, + arrays=arrays, + ) + + return metric_rows, curve_accumulator + + def evaluate_landmark_calibration( *, model: Any, @@ -660,6 +1171,7 @@ def evaluate_landmark_calibration( max_ipcw_weight: float, exclude_death_competing: bool, probability_bins: np.ndarray, + num_workers_calibration: int, ) -> tuple[pd.DataFrame, pd.DataFrame]: model.eval().to(device) hidden_all, row_arrays = infer_landmark_hidden( @@ -691,152 +1203,121 @@ def evaluate_landmark_calibration( death_index = int( getattr(model, "death_idx", getattr(model, "vocab_size", 1) - 1) ) + strata = _build_calibration_strata( + row_arrays["sex"], + row_arrays["landmark_age"], + ) + worker_count = max( + 1, + min( + int(num_workers_calibration), + len(disease_ids), + disease_chunk_size, + ), + ) + print(f"Calibration CPU workers: {worker_count}") 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, - } - ) + ] = defaultdict(_empty_curve_values) - 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, + executor = ( + ThreadPoolExecutor( + max_workers=worker_count, + thread_name_prefix="calibration", ) - - for column_index, token in enumerate( - tqdm( - chunk_ids, - desc=f"Calibration chunk {chunk_index}", - leave=False, - dynamic_ncols=True, - ) + if worker_count > 1 + else None + ) + try: + for chunk_index, chunk_ids in enumerate( + tqdm(chunks, desc="Disease chunks", 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)] + 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 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] + task_kwargs = [ + { + "column_index": column_index, + "token": int(token), + "logits_chunk": logits_chunk, + "rho_chunk": rho_chunk, + "strata": strata, + "row_patient_id": row_arrays["patient_id"], + "row_followup_end": row_arrays[ + "followup_end_time" + ], + "row_death_time": row_arrays["death_time"], + "first_occurrence_by_token": ( + landmark_dataset.first_occurrence_by_token + ), + "patient_count": patient_count, + "death_tokens": death_tokens, + "label_id_to_code": ( + landmark_dataset.dataset.label_id_to_code + ), + "dist_mode": dist_mode, + "horizons": horizons, + "death_index": death_index, + "min_cases": min_cases, + "min_controls": min_controls, + "max_ipcw_weight": max_ipcw_weight, + "exclude_death_competing": ( + exclude_death_competing + ), + "probability_bins": probability_bins, + } + for column_index, token in enumerate(chunk_ids) + ] + if executor is None: + chunk_results = ( + _evaluate_calibration_token(**kwargs) + for kwargs in tqdm( + task_kwargs, + desc=f"Calibration chunk {chunk_index}", + leave=False, + dynamic_ncols=True, + ) ) - 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), + for rows, local_curve in chunk_results: + metric_rows.extend(rows) + _merge_curve_accumulator( + curve_accumulator, + local_curve, ) - 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, + else: + futures = [ + executor.submit( + _evaluate_calibration_token, + **kwargs, + ) + for kwargs in task_kwargs + ] + for future in tqdm( + as_completed(futures), + total=len(futures), + desc=f"Calibration chunk {chunk_index}", + leave=False, + dynamic_ncols=True, + ): + rows, local_curve = future.result() + metric_rows.extend(rows) + _merge_curve_accumulator( + curve_accumulator, + local_curve, ) - 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 + del logits_chunk, rho_chunk + finally: + if executor is not None: + executor.shutdown(wait=True, cancel_futures=True) if not metric_rows: raise RuntimeError( @@ -844,7 +1325,14 @@ def evaluate_landmark_calibration( "horizons, min_cases, and disease selection." ) - metrics = pd.DataFrame(metric_rows) + metrics = ( + pd.DataFrame(metric_rows) + .sort_values( + ["token", "sex", "landmark_age", "horizon"], + kind="stable", + ) + .reset_index(drop=True) + ) curve_rows: List[Dict[str, Any]] = [] for (outcome, sex, horizon, bin_index), values in sorted( curve_accumulator.items() @@ -1331,6 +1819,13 @@ def main() -> None: disease_chunk_size = int( cfg_get(args, cfg, "disease_chunk_size", 64) ) + num_workers_calibration = int( + cfg_get(args, cfg, "num_workers_calibration", 0) + ) + if num_workers_calibration < 0: + raise ValueError("num_workers_calibration must be >= 0") + if num_workers_calibration == 0: + num_workers_calibration = max(1, int(os.cpu_count() or 1)) min_cases = int(cfg_get(args, cfg, "min_cases", 2)) min_controls = int(args.min_controls) if min_cases < 1: @@ -1385,6 +1880,7 @@ def main() -> None: max_ipcw_weight=float(args.max_ipcw_weight), exclude_death_competing=exclude_death_competing, probability_bins=probability_bins, + num_workers_calibration=num_workers_calibration, ) token_metrics = aggregate_metric_rows( landmark_metrics, @@ -1472,6 +1968,7 @@ def main() -> None: "horizons": [float(value) for value in horizons], "min_cases": min_cases, "min_controls": min_controls, + "num_workers_calibration": num_workers_calibration, "exclude_death_competing": exclude_death_competing, "landmark_estimand": ( "first-onset fixed-horizon risk; for non-death outcomes, " diff --git a/evaluate_calibration_all_runs_linux.sh b/evaluate_calibration_all_runs_linux.sh index 66ca2b3..a971ce6 100644 --- a/evaluate_calibration_all_runs_linux.sh +++ b/evaluate_calibration_all_runs_linux.sh @@ -29,6 +29,7 @@ LOG_ROOT="$SCRIPT_DIR/batch_logs/evaluate_calibration_all_runs" GPU_CSV="0" PYTHON_BIN="${PYTHON_BIN:-python}" NUM_WORKERS=4 +NUM_WORKERS_CALIBRATION=0 BATCH_SIZE=128 DISEASE_CHUNK_SIZE=64 HORIZONS="" @@ -52,6 +53,8 @@ Options: --python PATH Python executable (default: $PYTHON_BIN or python). --num-workers N DataLoader workers per job (default: 4). + --num-workers-calibration N CPU calibration workers per job. Default: 0, + which divides all logical CPUs across GPUs. --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. @@ -107,6 +110,14 @@ while (($# > 0)); do NUM_WORKERS="$2" shift 2 ;; + --num-workers-calibration) + [[ $# -ge 2 ]] || { + echo "ERROR: --num-workers-calibration requires a value." >&2 + exit 2 + } + NUM_WORKERS_CALIBRATION="$2" + shift 2 + ;; --batch-size) [[ $# -ge 2 ]] || { echo "ERROR: --batch-size requires a value." >&2 @@ -163,6 +174,10 @@ done echo "ERROR: --num-workers must be a non-negative integer." >&2 exit 2 } +[[ "$NUM_WORKERS_CALIBRATION" =~ ^[0-9]+$ ]] || { + echo "ERROR: --num-workers-calibration 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 @@ -203,6 +218,14 @@ for gpu in "${GPU_IDS[@]}"; do SEEN_GPUS["$gpu"]=1 done +if ((NUM_WORKERS_CALIBRATION == 0)); then + TOTAL_CPUS="$(nproc)" + NUM_WORKERS_CALIBRATION=$(( (TOTAL_CPUS + ${#GPU_IDS[@]} - 1) / ${#GPU_IDS[@]} )) + if ((NUM_WORKERS_CALIBRATION < 1)); then + NUM_WORKERS_CALIBRATION=1 + fi +fi + declare -a JOB_RUN_DIRS=() declare -a JOB_LOG_FILES=() @@ -281,6 +304,7 @@ run_job() { --eval_split test --device cuda --num_workers "$NUM_WORKERS" + --num_workers_calibration "$NUM_WORKERS_CALIBRATION" --batch_size "$BATCH_SIZE" --disease_chunk_size "$DISEASE_CHUNK_SIZE" ) @@ -341,6 +365,7 @@ worker() { echo "Runs root: $RUNS_ROOT" echo "GPUs: ${GPU_IDS[*]}" +echo "Calibration CPU workers per GPU job: $NUM_WORKERS_CALIBRATION" echo "Runs discovered: $run_count" echo "Incomplete runs skipped: $incomplete_count" echo "next_token runs skipped: $next_token_count" diff --git a/tests/test_calibration_metrics.py b/tests/test_calibration_metrics.py index 26fc110..8570345 100644 --- a/tests/test_calibration_metrics.py +++ b/tests/test_calibration_metrics.py @@ -1,17 +1,66 @@ import math import unittest +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import patch import numpy as np import pandas as pd +import torch from evaluate_calibration import ( + _censoring_km, + _evaluate_calibration_token, + _risk_probability_matrix, aggregate_metric_rows, compute_ipcw_cell, + compute_ipcw_horizons, + evaluate_landmark_calibration, fit_weighted_logistic_calibration, + fit_weighted_logistic_calibration_batch, ) +from evaluate_auc_v2 import _score_to_probability class IPCWCalibrationMetricTests(unittest.TestCase): + @staticmethod + def _naive_censoring_km(observed_times, censor_events): + observed_times = np.asarray(observed_times, dtype=np.float64) + censor_events = np.asarray(censor_events, dtype=bool) + event_times = np.unique(observed_times[censor_events]) + survival = 1.0 + survival_after = [] + for time_value in event_times: + at_risk = np.sum(observed_times >= time_value) + censored = np.sum( + censor_events & (observed_times == time_value) + ) + survival *= 1.0 - float(censored) / float(at_risk) + survival_after.append(survival) + return event_times, np.asarray(survival_after) + + def test_sorted_censoring_km_matches_naive_reference(self): + rng = np.random.RandomState(12) + observed_times = rng.randint(1, 20, size=500).astype(np.float64) + censor_events = rng.uniform(size=500) < 0.4 + + expected_times, expected_survival = self._naive_censoring_km( + observed_times, + censor_events, + ) + actual_times, actual_survival = _censoring_km( + observed_times, + censor_events, + ) + + np.testing.assert_array_equal(actual_times, expected_times) + np.testing.assert_allclose( + actual_survival, + expected_survival, + rtol=1e-14, + atol=1e-14, + ) + def test_no_censoring_matches_binary_metrics(self): result = compute_ipcw_cell( probabilities=np.asarray([0.2, 0.8]), @@ -83,6 +132,278 @@ class IPCWCalibrationMetricTests(unittest.TestCase): self.assertAlmostEqual(intercept, 0.0, places=7) self.assertAlmostEqual(slope, 1.0, places=7) + def test_batched_calibration_fits_multiple_horizons(self): + probabilities = np.vstack( + [ + np.repeat([0.1, 0.3, 0.7, 0.9], 100), + np.repeat([0.2, 0.4, 0.6, 0.8], 100), + ] + ) + outcomes = np.vstack( + [ + 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)], + ] + ), + np.concatenate( + [ + np.r_[np.ones(20), np.zeros(80)], + np.r_[np.ones(40), np.zeros(60)], + np.r_[np.ones(60), np.zeros(40)], + np.r_[np.ones(80), np.zeros(20)], + ] + ), + ] + ) + + calibration_in_large, intercept, slope = ( + fit_weighted_logistic_calibration_batch( + probabilities, + outcomes, + np.ones_like(probabilities), + ) + ) + + np.testing.assert_allclose(calibration_in_large, 0.0, atol=1e-7) + np.testing.assert_allclose(intercept, 0.0, atol=1e-7) + np.testing.assert_allclose(slope, 1.0, atol=1e-7) + + def test_all_horizons_reuse_one_censoring_km(self): + probabilities = np.asarray( + [ + [0.05, 0.10, 0.15, 0.20, 0.25], + [0.10, 0.20, 0.30, 0.40, 0.50], + [0.20, 0.35, 0.50, 0.65, 0.80], + ] + ) + event_times = np.asarray([0.5, 1.5, 4.0, np.inf, np.inf]) + censor_times = np.asarray([5.0, 5.0, 5.0, 2.5, 5.0]) + + with patch( + "evaluate_calibration._censoring_km", + wraps=_censoring_km, + ) as km: + results = compute_ipcw_horizons( + probabilities=probabilities, + event_times=event_times, + censor_times=censor_times, + horizons=np.asarray([1.0, 2.0, 5.0]), + min_cases=1, + min_controls=1, + max_ipcw_weight=0.0, + ) + + self.assertEqual(km.call_count, 1) + self.assertEqual(len(results), 3) + self.assertTrue(all(result is not None for result in results)) + self.assertEqual([result[0]["n_events"] for result in results], [1, 2, 3]) + + def test_batched_risk_probabilities_match_scalar_reference(self): + logits = np.asarray([-2.0, -0.5, 0.2, 1.5], dtype=np.float32) + rho = np.asarray([0.8, 1.0, 1.2, 1.5], dtype=np.float32) + horizons = np.asarray([0.1, 1.0, 5.0], dtype=np.float32) + + for dist_mode, token, death_idx, selected_rho in ( + ("exponential", 4, 9, None), + ("weibull", 4, 9, rho), + ("mixed", 9, 9, rho), + ("mixed", 4, 9, None), + ): + actual = _risk_probability_matrix( + logits=logits, + rho=selected_rho, + horizons=horizons, + dist_mode=dist_mode, + token=token, + death_idx=death_idx, + ) + expected = np.vstack( + [ + _score_to_probability( + logits, + selected_rho, + score_mode="risk", + horizon=float(horizon), + dist_mode=dist_mode, + token=token, + death_idx=death_idx, + ) + for horizon in horizons + ] + ) + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-7) + + def test_per_disease_worker_is_thread_safe(self): + logits_chunk = np.asarray( + [ + [-2.0, -1.5], + [-1.0, -0.5], + [0.0, 0.5], + [0.5, 1.0], + [1.0, 1.5], + [1.5, 2.0], + ], + dtype=np.float32, + ) + common = { + "logits_chunk": logits_chunk, + "rho_chunk": None, + "strata": [("Female", 50.0, np.arange(6, dtype=np.int64))], + "row_patient_id": np.arange(6, dtype=np.int32), + "row_followup_end": np.full(6, 65.0, dtype=np.float32), + "row_death_time": np.full(6, np.inf, dtype=np.float32), + "first_occurrence_by_token": { + 4: ( + np.asarray([0, 1], dtype=np.int32), + np.asarray([50.5, 52.0], dtype=np.float32), + ), + 5: ( + np.asarray([2, 3], dtype=np.int32), + np.asarray([50.7, 53.0], dtype=np.float32), + ), + }, + "patient_count": 6, + "death_tokens": set(), + "label_id_to_code": {4: "D4", 5: "D5"}, + "dist_mode": "exponential", + "horizons": np.asarray([1.0, 5.0], dtype=np.float32), + "death_index": 9, + "min_cases": 1, + "min_controls": 1, + "max_ipcw_weight": 0.0, + "exclude_death_competing": True, + "probability_bins": np.asarray([0.0, 0.5, 1.0]), + } + tasks = [ + {"column_index": 0, "token": 4, **common}, + {"column_index": 1, "token": 5, **common}, + ] + serial = [ + _evaluate_calibration_token(**task) + for task in tasks + ] + with ThreadPoolExecutor(max_workers=2) as executor: + parallel = list( + executor.map( + lambda task: _evaluate_calibration_token(**task), + tasks, + ) + ) + + for (serial_rows, serial_curve), ( + parallel_rows, + parallel_curve, + ) in zip(serial, parallel): + pd.testing.assert_frame_equal( + pd.DataFrame(serial_rows), + pd.DataFrame(parallel_rows), + ) + self.assertEqual(set(serial_curve), set(parallel_curve)) + for key in serial_curve: + self.assertEqual( + set(serial_curve[key]), + set(parallel_curve[key]), + ) + np.testing.assert_allclose( + list(serial_curve[key].values()), + list(parallel_curve[key].values()), + equal_nan=True, + ) + + def test_landmark_evaluation_parallel_matches_serial(self): + logits_chunk = np.asarray( + [ + [-2.0, -1.5], + [-1.0, -0.5], + [0.0, 0.5], + [0.5, 1.0], + [1.0, 1.5], + [1.5, 2.0], + ], + dtype=np.float32, + ) + row_arrays = { + "patient_id": np.arange(6, dtype=np.int32), + "sex": np.zeros(6, dtype=np.int8), + "landmark_age": np.full(6, 50.0, dtype=np.float32), + "followup_end_time": np.full(6, 65.0, dtype=np.float32), + "death_time": np.full(6, np.inf, dtype=np.float32), + } + landmark_dataset = SimpleNamespace( + subset_indices=np.arange(6, dtype=np.int64), + death_token_ids=[], + first_occurrence_by_token={ + 4: ( + np.asarray([0, 1], dtype=np.int32), + np.asarray([50.5, 52.0], dtype=np.float32), + ), + 5: ( + np.asarray([2, 3], dtype=np.int32), + np.asarray([50.7, 53.0], dtype=np.float32), + ), + }, + dataset=SimpleNamespace( + label_id_to_code={4: "D4", 5: "D5"} + ), + ) + + class FakeModel: + death_idx = 9 + vocab_size = 10 + + def eval(self): + return self + + def to(self, _device): + return self + + common = { + "model": FakeModel(), + "loader": [], + "landmark_dataset": landmark_dataset, + "disease_ids": [4, 5], + "dist_mode": "exponential", + "horizons": np.asarray([1.0, 5.0], dtype=np.float32), + "device": torch.device("cpu"), + "use_amp": False, + "hidden_cache_dtype": "float16", + "logit_batch_size": 8, + "disease_chunk_size": 2, + "min_cases": 1, + "min_controls": 1, + "max_ipcw_weight": 0.0, + "exclude_death_competing": True, + "probability_bins": np.asarray([0.0, 0.5, 1.0]), + } + with ( + patch( + "evaluate_calibration.infer_landmark_hidden", + return_value=( + np.zeros((6, 4), dtype=np.float16), + row_arrays, + ), + ), + patch( + "evaluate_calibration.project_distribution_chunk", + return_value=(logits_chunk, None), + ), + ): + serial_metrics, serial_curve = evaluate_landmark_calibration( + **common, + num_workers_calibration=1, + ) + parallel_metrics, parallel_curve = evaluate_landmark_calibration( + **common, + num_workers_calibration=2, + ) + + pd.testing.assert_frame_equal(serial_metrics, parallel_metrics) + pd.testing.assert_frame_equal(serial_curve, parallel_curve) + def test_metric_aggregation_uses_contribution_sums(self): metrics = pd.DataFrame( [