survival-model-toolkit

Shared helpers and analysis steps for a competing-risk-aware Cox survival modelling pipeline: a persisted train/test split (so multiple scripts never silently diverge), preprocessing fitted on training data only, discrimination and calibration metrics with bootstrap intervals, proportional-hazards diagnostics, descriptive/reporting tables, and a set of higher-level analysis steps for comparing models and validating design choices.

This package builds on two separately published, more narrowly scoped packages rather than duplicating their logic:

• shap-recommender -- exclusion / non-linearity / interaction recommendations from SHAP attributions.
• competing-risk-sensitivity -- Aalen-Johansen cumulative incidence and Fine-Gray export for a competing event such as death.

Install

pip install survival-model-toolkit

What's inside

Splitting and preprocessing

• make_or_load_split -- one train/test partition, persisted to disk, so every script that loads it sees an identical partition.
• temporal_split -- a temporal (index-date cutoff) split, for a sensitivity analysis closer to deployment than a random split.
• build_preprocessor -- one-hot or ordinal encoding, fitted on the training subset only; continuous variables are left on their natural scale by default so hazard ratios and SHAP attributions stay in interpretable units.
• NonlinearTransform -- quadratic or restricted-cubic-spline expansion, fit on train and replayed unchanged on test.
• add_interactions -- hierarchy-aware pairwise products: a pair whose main effect was dropped is skipped and logged, not silently omitted.
• onehot_group_map / aggregate_shap -- map a one-hot-encoded feature's dummy columns back to a single SHAP attribution.

Discrimination and calibration

• cox_risk_score / batch_risk_scores -- a scalar risk ordering for any fitted model exposing decision_function, predict, or predict_partial_hazard.
• MetricEval -- Harrell's C with a genuine bootstrap CI (every replicate resamples the observations, rather than rescoring a fixed test set), Uno's C, time-dependent AUC, Brier score / IBS, and a paired bootstrap for the difference between two models' C-index. Every model scored by one evaluator is bootstrapped on the same fixed set of resamples (drawn once from seed), so model comparisons are paired; each score's replicates are cached, and they run in parallel over n_jobs workers (default N_JOBS = 4) with per-replicate seeds, so results depend on seed alone, not on n_jobs.
• CalibrationPerform -- binned observed-vs-predicted calibration at a fixed horizon with bootstrap CIs per bin, the conventional calibration slope, an overlay plot across models, and an operational "is calibration stable" check (95% CI of the slope contains 1).

Diagnostics and descriptive tables

• ph_assumption_report -- global and per-term Schoenfeld residual tests.
• table1_with_smd -- baseline characteristics with standardised mean differences between groups.
• incidence_by_group -- crude incidence per 1,000 person-years with exact Poisson intervals.
• design_report -- design-matrix dimensionality and events-per-parameter.
• export_coefficients -- hazard ratios with CIs for a fitted scikit-survival Cox model (obtained via a matched lifelines refit, since scikit-survival itself has no covariance matrix).
• competing_risk_report / finegray_export -- thin wrappers around competing-risk-sensitivity.

Pipeline steps

• discovery_confirmation_split / joint_lrt -- screen candidate terms on one half of the data, confirm them jointly on the other.
• fit_and_score / comparator_models -- fit-and-report a Cox model, plus two useful comparators (restricted cubic splines; a penalised all-pairs interaction model).
• sequential_ablation -- C-index after each pipeline component, under every ordering.
• subgroup_performance -- discrimination within subgroups (e.g. for a fairness/equity audit).
• interaction_dose_response / plot_dose_response -- refit after adding the top-N interactions (ranked by effect size) for a grid of N, so the number admitted is chosen by held-out discrimination.
• margin_sensitivity_cindex -- regenerate recommendations at several subgroup margins and refit, reporting the margin's effect on discrimination.

Example

from survival_model_toolkit import (
    make_or_load_split, build_preprocessor, MetricEval, CalibrationPerform,
    cox_risk_score,
)
from sksurv.linear_model import CoxPHSurvivalAnalysis

X_train, X_test, y_train, y_test = make_or_load_split(X, y, path="split.json")

pre = build_preprocessor(onehot_cols=["sex"], contin_cols=["age"]).fit(X_train)
X_train_t, X_test_t = pre.transform(X_train), pre.transform(X_test)

model = CoxPHSurvivalAnalysis(alpha=1e-6, ties="efron").fit(X_train_t, y_train)

evaluator = MetricEval()
c, (lo, hi) = evaluator.boot_metric(y_test, cox_risk_score(model, X_test_t))
print(f"C-index = {c:.3f} (95% CI {lo:.3f}-{hi:.3f})")

calib = CalibrationPerform(t0=365.0)
print(calib.report(model, X_test_t, y_test, label="cox"))

License

MIT