Skills › Scientific Agent Skills › scikit-survival › train_survival_model.py
train_survival_model.py
586 lines · skills/scikit-survival/scripts/train_survival_model.py at 5f2bc03
1#!/usr/bin/env python3
2"""Train a leakage-safe Cox or ensemble example on explicit local schema."""
4from __future__ import annotations
6import argparse
7from importlib.metadata import version
8from typing import Any
10from _common import (
11 DEFAULT_SEED,
12 MAX_FEATURES,
13 CliError,
14 atomic_save_npz,
15 bounded_int,
16 emit_json,
17 parse_names,
18 probability,
19 read_csv,
20 structured_survival,
21 synthetic_survival_frame,
22)
25MODEL_CHOICES = (
26 "coxph",
27 "coxnet",
28 "random-forest",
29 "extra-trees",
30 "gradient-boosting",
31)
34def resolve_schema(
35 frame: Any,
36 *,
37 event_column: str,
38 time_column: str,
39 numeric_columns: list[str],
40 categorical_columns: list[str],
41) -> tuple[Any, Any, dict[str, Any]]:
42 """Select explicit features and build a validated structured outcome."""
44 if not numeric_columns and not categorical_columns:
45 raise CliError(
46 "an explicit schema is required: provide --numeric-columns and/or "
47 "--categorical-columns"
48 )
49 feature_columns = numeric_columns + categorical_columns
50 if len(feature_columns) > MAX_FEATURES:
51 raise CliError(f"at most {MAX_FEATURES} features are allowed")
52 if len(feature_columns) != len(set(feature_columns)):
53 raise CliError("numeric and categorical feature lists must be disjoint")
54 if event_column == time_column:
55 raise CliError("event and time columns must differ")
56 forbidden = {event_column, time_column}.intersection(feature_columns)
57 if forbidden:
58 raise CliError(
59 "outcomes must not be used as features: " + ", ".join(sorted(forbidden))
60 )
61 required = [event_column, time_column, *feature_columns]
62 missing = [name for name in required if name not in frame.columns]
63 if missing:
64 raise CliError(f"missing schema columns: {', '.join(missing)}")
66 y = structured_survival(
67 frame[event_column],
68 frame[time_column],
69 event_name=event_column,
70 time_name=time_column,
71 )
72 X = frame.loc[:, feature_columns].copy()
73 schema = {
74 "categorical_features": categorical_columns,
75 "event": event_column,
76 "numeric_features": numeric_columns,
77 "time": time_column,
78 }
79 return X, y, schema
82def build_pipeline(
83 model_name: str,
84 numeric_columns: list[str],
85 categorical_columns: list[str],
86 *,
87 seed: int,
88):
89 """Build a preprocessing-and-model pipeline without fitting it."""
91 try:
92 from sklearn.compose import ColumnTransformer
93 from sklearn.impute import SimpleImputer
94 from sklearn.pipeline import Pipeline
95 from sklearn.preprocessing import OneHotEncoder, StandardScaler
96 from sksurv.ensemble import (
97 ExtraSurvivalTrees,
98 GradientBoostingSurvivalAnalysis,
99 RandomSurvivalForest,
100 )
101 from sksurv.linear_model import (
102 CoxnetSurvivalAnalysis,
103 CoxPHSurvivalAnalysis,
104 )
105 except ImportError as exc:
106 from _common import PINNED_INSTALL
108 raise CliError(
109 f"survival stack unavailable; install with `{PINNED_INSTALL}`"
110 ) from exc
112 transformers: list[tuple[str, Any, list[str]]] = []
113 if numeric_columns:
114 numeric = Pipeline(
115 [
116 ("imputer", SimpleImputer(strategy="median")),
117 ("scale", StandardScaler()),
118 ]
119 )
120 transformers.append(("numeric", numeric, numeric_columns))
121 if categorical_columns:
122 categorical = Pipeline(
123 [
124 ("imputer", SimpleImputer(strategy="most_frequent")),
125 (
126 "encode",
127 OneHotEncoder(
128 drop="first",
129 handle_unknown="ignore",
130 sparse_output=False,
131 ),
132 ),
133 ]
134 )
135 transformers.append(("categorical", categorical, categorical_columns))
136 preprocess = ColumnTransformer(
137 transformers,
138 remainder="drop",
139 sparse_threshold=0.0,
140 verbose_feature_names_out=False,
141 )
143 if model_name == "coxph":
144 model = CoxPHSurvivalAnalysis(alpha=0.1, ties="efron")
145 elif model_name == "coxnet":
146 model = CoxnetSurvivalAnalysis(
147 alphas=[0.05],
148 l1_ratio=0.5,
149 fit_baseline_model=True,
150 )
151 elif model_name == "random-forest":
152 model = RandomSurvivalForest(
153 n_estimators=64,
154 min_samples_leaf=5,
155 n_jobs=1,
156 random_state=seed,
157 )
158 elif model_name == "extra-trees":
159 model = ExtraSurvivalTrees(
160 n_estimators=64,
161 min_samples_leaf=5,
162 n_jobs=1,
163 random_state=seed,
164 )
165 elif model_name == "gradient-boosting":
166 model = GradientBoostingSurvivalAnalysis(
167 n_estimators=64,
168 learning_rate=0.05,
169 max_depth=2,
170 random_state=seed,
171 )
172 else:
173 raise CliError(f"unsupported model: {model_name}")
174 return Pipeline([("preprocess", preprocess), ("model", model)])
177def parameter_grid(model_name: str) -> dict[str, list[Any]]:
178 """Return a deliberately small, bounded tuning grid."""
180 if model_name == "coxph":
181 return {"model__alpha": [0.01, 0.1, 1.0]}
182 if model_name == "coxnet":
183 return {"model__alphas": [[0.01], [0.05], [0.2]]}
184 if model_name in {"random-forest", "extra-trees"}:
185 return {"model__min_samples_leaf": [3, 6, 10]}
186 if model_name == "gradient-boosting":
187 return {
188 "model__learning_rate": [0.03, 0.1],
189 "model__max_depth": [1, 2],
190 }
191 raise CliError(f"unsupported model: {model_name}")
194def _event_field(y: Any) -> str:
195 return str(y.dtype.names[0])
198def _time_field(y: Any) -> str:
199 return str(y.dtype.names[1])
202def stratified_splits(y: Any, folds: int, *, seed: int) -> list[tuple[Any, Any]]:
203 """Create deterministic event-stratified indices."""
205 try:
206 import numpy as np
207 from sklearn.model_selection import StratifiedKFold
208 except ImportError as exc:
209 raise CliError("scikit-learn and NumPy are required") from exc
210 labels = y[_event_field(y)].astype(int)
211 counts = np.bincount(labels, minlength=2)
212 if counts.min() < folds:
213 raise CliError(
214 f"{folds}-fold CV requires at least {folds} events and censored rows"
215 )
216 splitter = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed)
217 return list(splitter.split(np.zeros(len(y)), labels))
220def uno_concordance(y_train: Any, y_test: Any, risk: Any) -> dict[str, Any]:
221 """Evaluate IPCW concordance using only the training censoring distribution."""
223 try:
224 import numpy as np
225 from sksurv.metrics import concordance_index_ipcw
226 except ImportError as exc:
227 raise CliError("scikit-survival and NumPy are required") from exc
228 time_train = y_train[_time_field(y_train)]
229 time_test = y_test[_time_field(y_test)]
230 tau = float(min(np.quantile(time_train, 0.8), np.quantile(time_test, 0.8)))
231 try:
232 score = float(concordance_index_ipcw(y_train, y_test, risk, tau=tau)[0])
233 except ValueError as exc:
234 return {
235 "available": False,
236 "reason": str(exc),
237 "training_censoring_distribution": True,
238 "truncation_time": tau,
239 }
240 return {
241 "available": True,
242 "score": score,
243 "training_censoring_distribution": True,
244 "truncation_time": tau,
245 }
248def nested_tune(
249 X: Any,
250 y: Any,
251 *,
252 model_name: str,
253 numeric_columns: list[str],
254 categorical_columns: list[str],
255 outer_folds: int,
256 inner_folds: int,
257 seed: int,
258) -> tuple[list[dict[str, Any]], dict[str, Any]]:
259 """Run nested CV on training data, then tune once on all training rows."""
261 try:
262 from sklearn.model_selection import GridSearchCV
263 from sksurv.metrics import concordance_index_censored
264 except ImportError as exc:
265 raise CliError("scikit-learn and scikit-survival are required") from exc
267 outer_results: list[dict[str, Any]] = []
268 for fold, (train_index, validation_index) in enumerate(
269 stratified_splits(y, outer_folds, seed=seed), start=1
270 ):
271 X_outer_train = X.iloc[train_index]
272 X_outer_validation = X.iloc[validation_index]
273 y_outer_train = y[train_index]
274 y_outer_validation = y[validation_index]
275 inner = stratified_splits(y_outer_train, inner_folds, seed=seed + fold)
276 search = GridSearchCV(
277 build_pipeline(
278 model_name,
279 numeric_columns,
280 categorical_columns,
281 seed=seed + fold,
282 ),
283 parameter_grid(model_name),
284 cv=inner,
285 error_score="raise",
286 n_jobs=1,
287 )
288 search.fit(X_outer_train, y_outer_train)
289 risk = search.predict(X_outer_validation)
290 harrell = float(
291 concordance_index_censored(
292 y_outer_validation[_event_field(y_outer_validation)],
293 y_outer_validation[_time_field(y_outer_validation)],
294 risk,
295 )[0]
296 )
297 outer_results.append(
298 {
299 "best_params": search.best_params_,
300 "fold": fold,
301 "harrell_c": harrell,
302 "uno_c": uno_concordance(y_outer_train, y_outer_validation, risk),
303 "validation_rows": int(len(validation_index)),
304 }
305 )
307 final_inner = stratified_splits(y, inner_folds, seed=seed + 10_000)
308 final_search = GridSearchCV(
309 build_pipeline(
310 model_name,
311 numeric_columns,
312 categorical_columns,
313 seed=seed,
314 ),
315 parameter_grid(model_name),
316 cv=final_inner,
317 error_score="raise",
318 n_jobs=1,
319 )
320 final_search.fit(X, y)
321 return outer_results, {
322 "best_estimator": final_search.best_estimator_,
323 "best_params": final_search.best_params_,
324 }
327def prediction_archive(
328 estimator: Any,
329 X_test: Any,
330 y_train: Any,
331 y_test: Any,
332 risk: Any,
333) -> dict[str, Any]:
334 """Build metric-evaluator input without serializing the estimator."""
336 import numpy as np
338 train_time = y_train[_time_field(y_train)]
339 test_time = y_test[_time_field(y_test)]
340 lower = float(np.quantile(test_time, 0.2))
341 upper = float(
342 min(
343 np.quantile(test_time, 0.8),
344 np.nextafter(train_time.max(), -np.inf),
345 )
346 )
347 if not lower < upper:
348 raise CliError("cannot construct a train-supported prediction time grid")
349 times = np.linspace(lower, upper, 8)
350 arrays: dict[str, Any] = {
351 "risk": np.asarray(risk, dtype=float),
352 "test_event": y_test[_event_field(y_test)],
353 "test_time": test_time,
354 "times": times,
355 "train_event": y_train[_event_field(y_train)],
356 "train_time": train_time,
357 }
358 if hasattr(estimator, "predict_survival_function"):
359 try:
360 functions = estimator.predict_survival_function(X_test)
361 arrays["survival"] = np.vstack([function(times) for function in functions])
362 except (AttributeError, NotImplementedError, ValueError) as exc:
363 raise CliError(f"could not export requested survival probabilities: {exc}") from exc
364 from evaluate_survival_metrics import validate_predictions
366 validate_predictions(arrays)
367 return arrays
370def train_and_report(
371 frame: Any,
372 *,
373 event_column: str,
374 time_column: str,
375 numeric_columns: list[str],
376 categorical_columns: list[str],
377 model_name: str,
378 test_fraction: float,
379 seed: int,
380 tune: bool,
381 outer_folds: int,
382 inner_folds: int,
383) -> tuple[Any, Any, Any, dict[str, Any]]:
384 """Split first, then fit all learned preprocessing inside a pipeline."""
386 try:
387 import numpy as np
388 from sklearn.model_selection import train_test_split
389 from sksurv.metrics import concordance_index_censored
390 except ImportError as exc:
391 raise CliError("the pinned survival stack is required") from exc
393 X, y, schema = resolve_schema(
394 frame,
395 event_column=event_column,
396 time_column=time_column,
397 numeric_columns=numeric_columns,
398 categorical_columns=categorical_columns,
399 )
400 X_train, X_test, y_train, y_test = train_test_split(
401 X,
402 y,
403 test_size=test_fraction,
404 random_state=seed,
405 stratify=y[event_column],
406 )
408 nested_results: list[dict[str, Any]] | None = None
409 best_params: dict[str, Any] | None = None
410 if tune:
411 nested_results, tuned = nested_tune(
412 X_train,
413 y_train,
414 model_name=model_name,
415 numeric_columns=numeric_columns,
416 categorical_columns=categorical_columns,
417 outer_folds=outer_folds,
418 inner_folds=inner_folds,
419 seed=seed,
420 )
421 estimator = tuned["best_estimator"]
422 best_params = tuned["best_params"]
423 else:
424 estimator = build_pipeline(
425 model_name, numeric_columns, categorical_columns, seed=seed
426 )
427 estimator.fit(X_train, y_train)
429 risk = estimator.predict(X_test)
430 harrell = float(
431 concordance_index_censored(y_test[event_column], y_test[time_column], risk)[0]
432 )
433 transformed_features = int(
434 estimator.named_steps["preprocess"].get_feature_names_out().shape[0]
435 )
436 nested_harrell = (
437 None
438 if not nested_results
439 else {
440 "mean": float(np.mean([item["harrell_c"] for item in nested_results])),
441 "standard_deviation": float(
442 np.std([item["harrell_c"] for item in nested_results], ddof=1)
443 )
444 if len(nested_results) > 1
445 else 0.0,
446 }
447 )
448 report = {
449 "data": {
450 "censored_test": int((~y_test[event_column]).sum()),
451 "censored_train": int((~y_train[event_column]).sum()),
452 "events_test": int(y_test[event_column].sum()),
453 "events_train": int(y_train[event_column].sum()),
454 "rows_test": int(len(y_test)),
455 "rows_train": int(len(y_train)),
456 },
457 "evaluation": {
458 "harrell_c": harrell,
459 "nested_cv_harrell_c": nested_harrell,
460 "uno_c": uno_concordance(y_train, y_test, risk),
461 },
462 "leakage_controls": {
463 "holdout_used_for_tuning": False,
464 "preprocessing_fit_scope": "training folds only",
465 "split_before_imputation_encoding_scaling": True,
466 "tuning_evaluation": "nested_cv_on_training_split" if tune else "not_run",
467 },
468 "model": {
469 "best_params": best_params,
470 "family": model_name,
471 "transformed_feature_count": transformed_features,
472 },
473 "package_versions": {
474 "numpy": version("numpy"),
475 "pandas": version("pandas"),
476 "scikit-learn": version("scikit-learn"),
477 "scikit-survival": version("scikit-survival"),
478 },
479 "schema": schema,
480 "schema_version": "1.0",
481 "seed": seed,
482 "tuning": {
483 "inner_folds": inner_folds if tune else None,
484 "outer_folds": outer_folds if tune else None,
485 "outer_results": nested_results,
486 "performed": tune,
487 },
488 "warning": (
489 "Predictive metrics do not establish calibration, transportability, "
490 "causal effects, or clinical utility. This tool is not clinical advice."
491 ),
492 }
493 return estimator, (X_train, y_train), (X_test, y_test, risk), report
496def build_parser() -> argparse.ArgumentParser:
497 parser = argparse.ArgumentParser(
498 description=(
499 "Train a leakage-safe Cox or ensemble example. Local CSV use requires "
500 "an explicit feature schema; no network access is performed."
501 )
502 )
503 parser.add_argument("--input", help="Optional bounded local .csv")
504 parser.add_argument("--event-column", default="event")
505 parser.add_argument("--time-column", default="time")
506 parser.add_argument("--numeric-columns", help="Comma-separated numeric features")
507 parser.add_argument(
508 "--categorical-columns", help="Comma-separated categorical features"
509 )
510 parser.add_argument("--model", choices=MODEL_CHOICES, default="coxph")
511 parser.add_argument("--test-fraction", type=probability, default=0.25)
512 parser.add_argument("--seed", type=bounded_int(0, 2**32 - 1), default=DEFAULT_SEED)
513 parser.add_argument(
514 "--synthetic-rows",
515 type=bounded_int(80, 20_000),
516 default=240,
517 )
518 parser.add_argument(
519 "--tune",
520 action="store_true",
521 help="Run bounded nested CV on training data before final holdout evaluation",
522 )
523 parser.add_argument("--outer-folds", type=bounded_int(2, 5), default=3)
524 parser.add_argument("--inner-folds", type=bounded_int(2, 5), default=3)
525 parser.add_argument(
526 "--prediction-output",
527 help="Optional local .npz for evaluate_survival_metrics.py",
528 )
529 parser.add_argument("--output", help="Optional local .json summary")
530 parser.add_argument("--force", action="store_true")
531 return parser
534def main(argv: list[str] | None = None) -> int:
535 parser = build_parser()
536 args = parser.parse_args(argv)
537 try:
538 if args.input:
539 frame, source_name = read_csv(args.input)
540 numeric = parse_names(args.numeric_columns)
541 categorical = parse_names(args.categorical_columns)
542 source = {"kind": "local_csv", "name": source_name}
543 else:
544 frame, default_numeric, default_categorical = synthetic_survival_frame(
545 rows=args.synthetic_rows, seed=args.seed
546 )
547 numeric = parse_names(args.numeric_columns) or default_numeric
548 categorical = parse_names(args.categorical_columns) or default_categorical
549 source = {
550 "kind": "synthetic",
551 "network_used": False,
552 "seed": args.seed,
553 }
555 estimator, train_data, test_data, report = train_and_report(
556 frame,
557 event_column=args.event_column,
558 time_column=args.time_column,
559 numeric_columns=numeric,
560 categorical_columns=categorical,
561 model_name=args.model,
562 test_fraction=args.test_fraction,
563 seed=args.seed,
564 tune=args.tune,
565 outer_folds=args.outer_folds,
566 inner_folds=args.inner_folds,
567 )
568 report["source"] = source
569 if args.prediction_output:
570 X_train, y_train = train_data
571 X_test, y_test, risk = test_data
572 del X_train
573 atomic_save_npz(
574 prediction_archive(estimator, X_test, y_train, y_test, risk),
575 args.prediction_output,
576 force=args.force,
577 )
578 report["prediction_output"] = args.prediction_output
579 emit_json(report, output=args.output, force=args.force)
580 except CliError as exc:
581 parser.error(str(exc))
582 return 0
585if __name__ == "__main__":
586 raise SystemExit(main())