Add all-future calibration evaluation
This commit is contained in:
1517
evaluate_calibration.py
Normal file
1517
evaluate_calibration.py
Normal file
File diff suppressed because it is too large
Load Diff
379
evaluate_calibration_all_runs_linux.sh
Normal file
379
evaluate_calibration_all_runs_linux.sh
Normal file
@@ -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
|
||||
143
tests/test_calibration_metrics.py
Normal file
143
tests/test_calibration_metrics.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user