#!/usr/bin/env python3 """Generate a survival modelling script: Cox, Weibull AFT, random survival forest and gradient boosting, compared with survival-appropriate metrics on a held-out set. The script it writes: - loads the CSV, one-hot encodes categorical covariates (drop first), splits train/test stratified on the event indicator (or by --time-col order for out-of-time evaluation) - fits lifelines CoxPHFitter (penalizer chosen from a small grid by 5-fold CV concordance), lifelines WeibullAFTFitter, and scikit-survival RandomSurvivalForest and GradientBoostingSurvivalAnalysis - evaluates on the test set: Harrell's C, Uno's C (IPCW, truncated at tau), time-dependent cumulative/ dynamic AUC at --times, integrated Brier score over the --times range (lower is better, 0.25 = useless for S near 0.5), and calibration at the middle time: predicted vs Kaplan-Meier observed survival by risk quintile - writes survival_model_comparison.csv, test_survival_predictions.csv (S(t) per test row and time) and the best model by integrated Brier score with joblib Needs lifelines, scikit-survival, scikit-learn, pandas, numpy, joblib. Usage: python gen_survival_script.py --data survival.csv --duration-col duration_days --event-col event --covariates plan age -o surv.py python gen_survival_script.py --data data.csv --duration-col t --event-col died --covariates age stage sex --times 90 180 365 python gen_survival_script.py --data data.csv --duration-col t --event-col churn --covariates plan --time-col signup_date Exit codes: 0 written, 2 usage error. """ import argparse import os import sys TEMPLATE = '''"""Survival models for {duration!r} / {event!r}. Generated by gen_survival_script.py. Run: python {script} """ import joblib import numpy as np import pandas as pd from lifelines import CoxPHFitter, KaplanMeierFitter, WeibullAFTFitter from lifelines.utils import k_fold_cross_validation from sklearn.model_selection import train_test_split from sksurv.ensemble import GradientBoostingSurvivalAnalysis, RandomSurvivalForest from sksurv.metrics import (concordance_index_censored, concordance_index_ipcw, cumulative_dynamic_auc, integrated_brier_score) from sksurv.util import Surv DATA, DURATION, EVENT, COVARIATES, TIME_COL = {data!r}, {duration!r}, {event!r}, {covariates!r}, {time_col!r} TIMES, TEST_SIZE, SEED = {times!r}, {test_size}, 0 df = pd.read_csv(DATA) df = df.dropna(subset=[DURATION, EVENT] + COVARIATES) df = df[df[DURATION] > 0] X = pd.get_dummies(df[COVARIATES], drop_first=True, dtype=float) data = pd.concat([X, df[[DURATION, EVENT]]], axis=1) if TIME_COL: order = df[TIME_COL].sort_values().index cut = int(len(order) * (1 - TEST_SIZE)) train, test = data.loc[order[:cut]], data.loc[order[cut:]] else: train, test = train_test_split(data, test_size=TEST_SIZE, stratify=data[EVENT], random_state=SEED) feats = list(X.columns) y_tr = Surv.from_arrays(train[EVENT].astype(bool), train[DURATION]) y_te = Surv.from_arrays(test[EVENT].astype(bool), test[DURATION]) if TIMES: times = np.array(TIMES, dtype=float) else: # inside the test follow-up, where IPCW metrics are defined times = np.quantile(test.loc[test[EVENT] == 1, DURATION], [0.2, 0.4, 0.6, 0.8]) times = times[(times > test[DURATION].min()) & (times < test[DURATION].max()) & (times < train[DURATION].max())] tau = float(times.max()) def cox_factory(): best = None for pen in (0.0, 0.01, 0.1, 1.0): scores = k_fold_cross_validation(CoxPHFitter(penalizer=pen), train, DURATION, EVENT, k=5, scoring_method="concordance_index", seed=SEED) if best is None or np.mean(scores) > best[0]: best = (np.mean(scores), pen) print(f"Cox penalizer {{best[1]}} (CV concordance {{best[0]:.3f}})") return CoxPHFitter(penalizer=best[1]).fit(train, DURATION, EVENT) models = {{}} models["cox"] = cox_factory() models["weibull_aft"] = WeibullAFTFitter(penalizer=0.01).fit(train, DURATION, EVENT) models["random_survival_forest"] = RandomSurvivalForest(n_estimators=300, min_samples_leaf=10, max_features="sqrt", random_state=SEED, n_jobs=-1).fit(train[feats], y_tr) models["gradient_boosting"] = GradientBoostingSurvivalAnalysis(n_estimators=200, learning_rate=0.05, max_depth=2, subsample=0.8, random_state=SEED).fit(train[feats], y_tr) def surv_matrix(name, m, frame): """Rows: test subjects, columns: times -> S(t).""" if name in ("cox", "weibull_aft"): sf = m.predict_survival_function(frame[feats], times=times) return sf.T.to_numpy() fns = m.predict_survival_function(frame[feats]) return np.vstack([fn(times) for fn in fns]) def risk_score(name, m, frame): if name == "cox": return m.predict_partial_hazard(frame[feats]).to_numpy() if name == "weibull_aft": return -m.predict_median(frame[feats]).to_numpy() # shorter median = higher risk return m.predict(frame[feats]) rows, preds = [], [] mid = times[len(times) // 2] for name, m in models.items(): S = surv_matrix(name, m, test) risk = risk_score(name, m, test) harrell = concordance_index_censored(test[EVENT].astype(bool), test[DURATION], risk)[0] uno = concordance_index_ipcw(y_tr, y_te, risk, tau=tau)[0] auc_t, mean_auc = cumulative_dynamic_auc(y_tr, y_te, risk, times) ibs = integrated_brier_score(y_tr, y_te, S, times) if len(times) > 1 else float("nan") rows.append({{"model": name, "harrell_c": harrell, "uno_c": uno, "mean_auc": mean_auc, "ibs": ibs, **{{f"auc_t{{t:g}}": a for t, a in zip(times, auc_t)}}}}) k = int(np.argmin(np.abs(times - mid))) q = pd.qcut(pd.Series(S[:, k]).rank(method="first"), 5, labels=False) cal = [] for g in range(5): idx = np.where(q.to_numpy() == g)[0] km = KaplanMeierFitter().fit(test[DURATION].iloc[idx], test[EVENT].iloc[idx]) cal.append((float(S[idx, k].mean()), float(km.survival_function_at_times(times[k]).iloc[0]))) rows[-1]["calibration_max_gap"] = max(abs(p - o) for p, o in cal) for i in range(len(test)): preds.append({{"model": name, "row": int(test.index[i]), **{{f"S_{{t:g}}": S[i, j] for j, t in enumerate(times)}}}}) table = pd.DataFrame(rows).sort_values("ibs") table.to_csv("survival_model_comparison.csv", index=False) pd.DataFrame(preds).to_csv("test_survival_predictions.csv", index=False) print(f"test rows: {{len(test)}}, events: {{int(test[EVENT].sum())}}, times: {{', '.join(f'{{t:g}}' for t in times)}}") print(table.round(4).to_string(index=False)) best = table.iloc[0]["model"] joblib.dump(models[best], f"best_survival_model_{{best}}.joblib") print(f"best by integrated Brier score: {{best}} -> best_survival_model_{{best}}.joblib") cox = models["cox"] try: from lifelines.statistics import proportional_hazard_test ph = proportional_hazard_test(cox, train, time_transform="rank").summary bad = ph[ph["p"] < 0.05] if len(bad): print("Cox PH assumption violated for:", ", ".join(map(str, bad.index))) except Exception as exc: # diagnostics only; never block the comparison print(f"PH test skipped: {{exc}}") ''' def main(): ap = argparse.ArgumentParser(description="Generate a survival model comparison script.", epilog=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--data", required=True) ap.add_argument("--duration-col", required=True) ap.add_argument("--event-col", required=True) ap.add_argument("--covariates", nargs="+", required=True) ap.add_argument("--times", type=float, nargs="*", help="evaluation times (default: event-time quantiles)") ap.add_argument("--time-col", help="order column for an out-of-time test split") ap.add_argument("--test-size", type=float, default=0.3) ap.add_argument("-o", "--output") a = ap.parse_args() if not 0.1 <= a.test_size <= 0.5: print("error: --test-size must be between 0.1 and 0.5", file=sys.stderr) return 2 text = TEMPLATE.format(data=a.data, duration=a.duration_col, event=a.event_col, covariates=a.covariates, time_col=a.time_col, times=a.times or [], test_size=a.test_size, script=os.path.basename(a.output or "survival_models.py")) if a.output: with open(a.output, "w", encoding="utf-8") as f: f.write(text) print(f"wrote {a.output}: Cox, Weibull AFT, RSF, GBSA on {', '.join(a.covariates)}", file=sys.stderr) else: sys.stdout.write(text) return 0 if __name__ == "__main__": sys.exit(main())