Skills › Scientific Agent Skills › scikit-survival › evaluate_survival_metrics.py
evaluate_survival_metrics.py
322 lines · skills/scikit-survival/scripts/evaluate_survival_metrics.py at 5f2bc03
1#!/usr/bin/env python3
2"""Evaluate censoring-aware survival metrics with strict input contracts."""
4from __future__ import annotations
6import argparse
7import zipfile
8from typing import Any
10from _common import (
11 DEFAULT_SEED,
12 MAX_INPUT_BYTES,
13 MAX_ROWS,
14 MAX_TIME_POINTS,
15 CliError,
16 bounded_int,
17 checked_input_file,
18 emit_json,
19 structured_survival,
20)
23REQUIRED_ARRAYS = {
24 "risk",
25 "test_event",
26 "test_time",
27 "times",
28 "train_event",
29 "train_time",
30}
33def synthetic_metric_inputs(seed: int = DEFAULT_SEED) -> dict[str, Any]:
34 """Create deterministic, non-clinical predictions with supported follow-up."""
36 try:
37 import numpy as np
38 except ImportError as exc:
39 raise CliError("NumPy is required") from exc
40 rng = np.random.default_rng(seed)
42 def observed(rows: int) -> tuple[Any, Any, Any]:
43 risk = rng.normal(size=rows)
44 event_time = rng.exponential(scale=9.0 * np.exp(-0.65 * risk)) + 0.1
45 censor_time = rng.exponential(scale=13.0, size=rows) + 0.1
46 event = event_time <= censor_time
47 time = np.minimum(event_time, censor_time)
48 return event, time, risk
50 train_event, train_time, _ = observed(220)
51 test_event, test_time, risk = observed(90)
52 support_limit = float(test_time.max() + 5.0)
53 train_event[-1] = False
54 train_time[-1] = max(support_limit, float(train_time.max()))
55 upper = float(min(np.quantile(test_time, 0.8), train_time.max() - 1e-6))
56 lower = float(np.quantile(test_time, 0.2))
57 times = np.linspace(lower, upper, 8)
58 survival = np.exp(-np.exp(0.65 * risk[:, None]) * times[None, :] / 9.0)
59 return {
60 "risk": risk,
61 "survival": survival,
62 "test_event": test_event,
63 "test_time": test_time,
64 "times": times,
65 "train_event": train_event,
66 "train_time": train_time,
67 }
70def load_prediction_archive(value: str) -> tuple[dict[str, Any], str]:
71 """Load a bounded NPZ archive without pickle."""
73 try:
74 import numpy as np
75 except ImportError as exc:
76 raise CliError("NumPy is required") from exc
77 path = checked_input_file(value, suffixes={".npz"})
78 try:
79 with zipfile.ZipFile(path) as archive:
80 members = archive.infolist()
81 if len(members) > len(REQUIRED_ARRAYS) + 1 or sum(
82 member.file_size for member in members
83 ) > MAX_INPUT_BYTES:
84 raise CliError("prediction archive exceeds expanded array limits")
85 if len({member.filename for member in members}) != len(members):
86 raise CliError("prediction archive contains duplicate members")
87 with np.load(path, allow_pickle=False) as archive:
88 names = set(archive.files)
89 missing = REQUIRED_ARRAYS - names
90 if missing:
91 raise CliError(
92 "prediction archive is missing: " + ", ".join(sorted(missing))
93 )
94 unexpected = names - REQUIRED_ARRAYS - {"survival"}
95 if unexpected:
96 raise CliError(
97 "prediction archive has unsupported arrays: "
98 + ", ".join(sorted(unexpected))
99 )
100 arrays = {name: archive[name] for name in archive.files}
101 except (OSError, ValueError, zipfile.BadZipFile) as exc:
102 raise CliError(f"cannot load safe NPZ archive {path.name}: {exc}") from exc
103 return arrays, path.name
106def validate_time_grid(
107 times: Any,
108 y_train: Any,
109 y_test: Any,
110) -> Any:
111 """Validate shape, order, observed support, and censoring support."""
113 try:
114 import numpy as np
115 from sksurv.nonparametric import CensoringDistributionEstimator
116 except ImportError as exc:
117 raise CliError("the pinned survival stack is required") from exc
118 values = np.asarray(times, dtype=float)
119 if values.ndim != 1:
120 raise CliError("times must be one-dimensional")
121 if not 2 <= values.size <= MAX_TIME_POINTS:
122 raise CliError(f"times must contain 2 through {MAX_TIME_POINTS} values")
123 if not np.isfinite(values).all() or (values <= 0).any():
124 raise CliError("times must be finite and strictly positive")
125 if not (np.diff(values) > 0).all():
126 raise CliError("times must be unique and strictly increasing")
128 train_time = y_train[y_train.dtype.names[1]]
129 test_time = y_test[y_test.dtype.names[1]]
130 if not test_time.max() < train_time.max():
131 raise CliError(
132 "test follow-up must end before training follow-up for IPCW support"
133 )
134 if values[0] <= test_time.min() or values[-1] >= test_time.max():
135 raise CliError(
136 "evaluation times must lie strictly inside the test follow-up range"
137 )
138 if values[-1] >= train_time.max():
139 raise CliError("evaluation times exceed training follow-up support")
140 censoring = CensoringDistributionEstimator().fit(y_train)
141 try:
142 probability = censoring.predict_proba(values)
143 except ValueError as exc:
144 raise CliError(
145 f"censoring distribution is undefined on the grid: {exc}"
146 ) from exc
147 if not np.isfinite(probability).all() or (probability <= 0).any():
148 raise CliError(
149 "training censoring survival must remain positive across the time grid"
150 )
151 try:
152 censoring.predict_ipcw(y_test)
153 except ValueError as exc:
154 raise CliError(f"test event times lack positive censoring support: {exc}") from exc
155 test_event = y_test[y_test.dtype.names[0]]
156 if not (test_event & (test_time <= values[0])).any():
157 raise CliError("dynamic AUC needs an observed case by the first evaluation time")
158 if not (test_time > values[-1]).any():
159 raise CliError("dynamic AUC needs controls after the final evaluation time")
160 return values
163def validate_predictions(
164 arrays: dict[str, Any],
165) -> tuple[Any, Any | None, Any, Any, Any]:
166 """Validate risk/survival shapes and return metric-ready arrays."""
168 try:
169 import numpy as np
170 except ImportError as exc:
171 raise CliError("NumPy is required") from exc
172 for name in ("train_event", "train_time", "test_event", "test_time"):
173 value = np.asarray(arrays[name])
174 if value.ndim != 1 or not 2 <= value.size <= MAX_ROWS:
175 raise CliError(f"{name} must have 2 through {MAX_ROWS} one-dimensional values")
176 y_train = structured_survival(arrays["train_event"], arrays["train_time"])
177 y_test = structured_survival(arrays["test_event"], arrays["test_time"])
178 times = validate_time_grid(arrays["times"], y_train, y_test)
179 n_test = len(y_test)
181 risk = np.asarray(arrays["risk"], dtype=float)
182 valid_risk_shape = risk.shape == (n_test,) or risk.shape == (
183 n_test,
184 len(times),
185 )
186 if not valid_risk_shape:
187 raise CliError("risk must have shape (n_test,) or (n_test, n_times)")
188 if not np.isfinite(risk).all():
189 raise CliError("risk contains non-finite values")
191 survival: Any | None = None
192 if "survival" in arrays:
193 survival = np.asarray(arrays["survival"], dtype=float)
194 if survival.shape != (n_test, len(times)):
195 raise CliError("survival must have shape (n_test, n_times)")
196 if not np.isfinite(survival).all():
197 raise CliError("survival contains non-finite values")
198 if ((survival < 0) | (survival > 1)).any():
199 raise CliError("survival probabilities must be within [0, 1]")
200 if (np.diff(survival, axis=1) > 1e-10).any():
201 raise CliError(
202 "each survival-probability row must be non-increasing over time"
203 )
204 return risk, survival, times, y_train, y_test
207def evaluate(arrays: dict[str, Any]) -> dict[str, Any]:
208 """Compute discrimination and prediction-error metrics by input type."""
210 try:
211 import numpy as np
212 from sksurv.metrics import (
213 brier_score,
214 concordance_index_censored,
215 concordance_index_ipcw,
216 cumulative_dynamic_auc,
217 integrated_brier_score,
218 )
219 except ImportError as exc:
220 raise CliError("the pinned survival stack is required") from exc
221 risk, survival, times, y_train, y_test = validate_predictions(arrays)
222 event_field, time_field = y_test.dtype.names
224 auc, mean_auc = cumulative_dynamic_auc(y_train, y_test, risk, times)
225 metrics: dict[str, Any] = {
226 "cumulative_dynamic_auc": {
227 "mean": float(mean_auc),
228 "times": times.tolist(),
229 "values": np.asarray(auc, dtype=float).tolist(),
230 }
231 }
232 if risk.ndim == 1:
233 metrics["harrell_c"] = float(
234 concordance_index_censored(y_test[event_field], y_test[time_field], risk)[0]
235 )
236 metrics["uno_c"] = float(
237 concordance_index_ipcw(y_train, y_test, risk, tau=float(times[-1]))[0]
238 )
239 else:
240 metrics["harrell_c"] = None
241 metrics["uno_c"] = None
242 metrics["concordance_note"] = (
243 "Skipped: concordance functions require one risk score per row; "
244 "the supplied risk is time-dependent."
245 )
247 if survival is not None:
248 returned_times, scores = brier_score(y_train, y_test, survival, times)
249 metrics["brier_score"] = {
250 "times": np.asarray(returned_times, dtype=float).tolist(),
251 "values": np.asarray(scores, dtype=float).tolist(),
252 }
253 metrics["integrated_brier_score"] = float(
254 integrated_brier_score(y_train, y_test, survival, times)
255 )
256 else:
257 metrics["brier_score"] = None
258 metrics["integrated_brier_score"] = None
259 metrics["brier_note"] = (
260 "Skipped: Brier metrics require survival probabilities, not risk scores."
261 )
263 if not np.isfinite(auc).all() or not np.isfinite(mean_auc):
264 raise CliError("dynamic AUC is undefined for these cases, controls, or weights")
265 return {
266 "assumptions": {
267 "censoring_distribution_fit_on_training_only": True,
268 "independent_censoring_from_features": "required_by_IPCW_estimator",
269 "test_follow_up_within_training_support": True,
270 "time_grid_within_test_follow_up": True,
271 },
272 "input_contract": {
273 "risk": ("higher means greater event risk; 1D or n_test-by-n_times"),
274 "survival": ("optional n_test-by-n_times probabilities; never risk scores"),
275 },
276 "metrics": metrics,
277 "schema_version": "1.0",
278 "warning": (
279 "Discrimination and prediction error do not establish calibration "
280 "at every horizon, causal effects, or clinical utility."
281 ),
282 }
285def build_parser() -> argparse.ArgumentParser:
286 parser = argparse.ArgumentParser(
287 description=(
288 "Evaluate risk scores and optional survival probabilities from a "
289 "bounded local NPZ archive. No input uses deterministic synthetic data."
290 )
291 )
292 parser.add_argument("--input", help="Optional local .npz prediction archive")
293 parser.add_argument("--seed", type=bounded_int(0, 2**32 - 1), default=DEFAULT_SEED)
294 parser.add_argument("--output", help="Optional local .json report")
295 parser.add_argument("--force", action="store_true")
296 return parser
299def main(argv: list[str] | None = None) -> int:
300 parser = build_parser()
301 args = parser.parse_args(argv)
302 try:
303 if args.input:
304 arrays, name = load_prediction_archive(args.input)
305 source = {"kind": "local_npz", "name": name}
306 else:
307 arrays = synthetic_metric_inputs(args.seed)
308 source = {
309 "kind": "synthetic",
310 "network_used": False,
311 "seed": args.seed,
312 }
313 report = evaluate(arrays)
314 report["source"] = source
315 emit_json(report, output=args.output, force=args.force)
316 except CliError as exc:
317 parser.error(str(exc))
318 return 0
321if __name__ == "__main__":
322 raise SystemExit(main())