# Machine learning
The Cox model and its parametric cousins assume a particular shape for how covariates act on the hazard (usually linear and proportional). When you have many covariates, non-linear effects, or interactions you did not anticipate, tree-based machine-learning models can predict better without those assumptions. Greenwood ships three, all fully non-parametric and all producing per-subject survival and cumulative-hazard *functions* rather than a single risk score: a [SurvivalTree](../reference/SurvivalTree.md#greenwood.SurvivalTree), the bagged [RandomSurvivalForest](../reference/RandomSurvivalForest.md#greenwood.RandomSurvivalForest) (and its extremely-randomized [ExtraSurvivalTrees](../reference/ExtraSurvivalTrees.md#greenwood.ExtraSurvivalTrees) variant), and [GradientBoostingSurvivalAnalysis](../reference/GradientBoostingSurvivalAnalysis.md#greenwood.GradientBoostingSurvivalAnalysis). This page introduces each, then shows how to score, interpret, and visualize them.
We'll work with `lung`-dataset outcomes, keeping the rows with complete covariates so the examples are easy to follow.
``` python
import greenwood as gw
# Load the lung dataset and keep rows with complete covariates
lung = gw.load_dataset("lung", backend="polars").drop_nulls(
subset=["ph.ecog", "ph.karno", "wt.loss"]
)
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))
covariates = ["age", "sex", "ph.ecog", "ph.karno", "wt.loss"]
y
```
Surv(type=right, n=213, events=151)
# Survival trees
A [SurvivalTree](../reference/SurvivalTree.md#greenwood.SurvivalTree) recursively splits subjects into groups whose survival curves differ as much as possible, measured by the log-rank statistic at each candidate split. Each leaf stores a Kaplan-Meier survival curve and a Nelson-Aalen cumulative hazard estimated from the subjects that land in it. A single tree is easy to reason about but high-variance; it is most useful as the interpretable building block of an ensemble.
``` python
# Grow a shallow, interpretable survival tree
tree = gw.SurvivalTree(max_depth=3, random_state=0).fit(y, lung[covariates])
tree
```
SurvivalTree (log-rank splits, 15 nodes, 8 leaves)
n = 213, events = 151, features = 5
Calling `predict` with the default `type="risk"` returns Ishwaran's *ensemble mortality* (the sum of the predicted cumulative hazard over the training event times). Larger values mean higher risk.
``` python
# Risk score for the first eight subjects (higher = higher risk)
tree.predict(lung[covariates])[:8]
```
array([ 95.15345288, 95.15345288, 95.15345288, 95.15345288,
95.15345288, 168.54325397, 168.54325397, 95.15345288])
# Random survival forests
A [RandomSurvivalForest](../reference/RandomSurvivalForest.md#greenwood.RandomSurvivalForest) averages many decorrelated survival trees, each grown on a bootstrap sample of the data and considering a random subset of covariates at every split. Averaging drives down the variance of the single tree and yields a well-calibrated, flexible model. Passing `oob_score=True` computes an out-of-bag concordance estimate (an honest, held-out measure of discrimination) using only the trees for which each subject was not in the bootstrap sample.
``` python
# Fit a forest and report the out-of-bag concordance
rsf = gw.RandomSurvivalForest(
n_estimators=200, oob_score=True, random_state=0
).fit(y, lung[covariates])
rsf
```
RandomSurvivalForest (200 log-rank trees, max_features='sqrt')
n = 213, events = 151, features = 5
out-of-bag concordance = 0.6061
## Predictions
Like the tree, the forest predicts a scalar risk score by default. It also predicts full survival or cumulative-hazard curves: pass `type="survival"` (or `"cumulative_hazard"`) and optional `times`, and you get one column per subject.
``` python
# Predicted survival probabilities at one and two years for three subjects
rsf.predict(lung[covariates][:3], type="survival", times=[365, 730])
```
PolarsRows2Columns4
|
|
time
f64
|
subject_1
f64
|
subject_2
f64
|
subject_3
f64
|
| 0 |
365 |
0.570687942169 |
0.664342402851 |
0.433306613867 |
| 1 |
730 |
0.0856037841238 |
0.366642713531 |
0.106499220759 |
## Variable importance
[variable_importance()](../reference/GradientBoostingSurvivalAnalysis.md#greenwood.GradientBoostingSurvivalAnalysis.variable_importance) measures how much each covariate contributes by permuting its values in the out-of-bag samples and recording the drop in concordance. Larger values indicate more important covariates. The same table is what [tidy()](../reference/tidy.md#greenwood.tidy) returns for a fitted forest.
``` python
# Permutation variable importance (out-of-bag)
rsf.variable_importance()
```
PolarsRows5Columns2
|
|
term
str
|
importance
f64
|
| 0 |
ph.ecog |
0.0552117647059 |
| 1 |
sex |
0.0398823529412 |
| 2 |
age |
0.00589411764706 |
| 3 |
ph.karno |
-0.00187058823529 |
| 4 |
wt.loss |
-0.00203529411765 |
# Extra survival trees
[ExtraSurvivalTrees](../reference/ExtraSurvivalTrees.md#greenwood.ExtraSurvivalTrees) is an extremely-randomized variant: in addition to sampling covariates, it draws the split *threshold* at random for each candidate rather than optimizing it. The extra randomness further decorrelates the trees, often lowering variance and speeding up fitting. Following the extra-trees convention, it grows each tree on the whole sample by default; set `bootstrap=True` when you want out-of-bag scoring or variable importance.
``` python
# Extremely-randomized survival trees with out-of-bag scoring enabled
ext = gw.ExtraSurvivalTrees(
n_estimators=200, bootstrap=True, oob_score=True, random_state=0
).fit(y, lung[covariates])
ext
```
ExtraSurvivalTrees (200 log-rank trees, max_features='sqrt')
n = 213, events = 151, features = 5
out-of-bag concordance = 0.6155
> **Tip: Accelerating the split search on large data**
>
> The best-split search is the dominant cost when fitting forests on large datasets. If you install the `fast` extra (`pip install greenwood[fast]`), you can pass `engine="numba"` to [RandomSurvivalForest](../reference/RandomSurvivalForest.md#greenwood.RandomSurvivalForest) or [ExtraSurvivalTrees](../reference/ExtraSurvivalTrees.md#greenwood.ExtraSurvivalTrees) for a compiled split kernel that avoids materializing the large risk-set matrices. The default `engine="numpy"` is deterministic and needs no extra dependency; the Numba path produces a statistically equivalent fit that may differ slightly in tie-breaking.
# Gradient-boosted survival
[GradientBoostingSurvivalAnalysis](../reference/GradientBoostingSurvivalAnalysis.md#greenwood.GradientBoostingSurvivalAnalysis) takes a different route. Instead of averaging independent trees, it builds an additive model by fitting each new (shallow) regression tree to the negative gradient of the Cox partial-likelihood loss (the martingale residuals under the current model) and adding a shrunken version of it. The result is a non-linear generalization of the Cox model that often predicts very well on structured tabular data.
``` python
# Fit a gradient-boosted survival model
gbm = gw.GradientBoostingSurvivalAnalysis(
n_estimators=300, learning_rate=0.05, max_depth=2, random_state=0
).fit(y, lung[covariates])
gbm
```
GradientBoostingSurvivalAnalysis (300 trees, learning_rate=0.05, max_depth=2)
n = 213, events = 151, features = 5
The additive score `F(x)` is the boosted analogue of the Cox linear predictor. `predict(type="lp")` returns it, `type="risk"` returns `exp(F(x))`, and `type="survival"` combines it with a Breslow baseline hazard to produce survival curves.
``` python
# Survival curves at one and two years for three subjects
gbm.predict(lung[covariates][:3], type="survival", times=[365, 730])
```
PolarsRows2Columns4
|
|
time
f64
|
subject_1
f64
|
subject_2
f64
|
subject_3
f64
|
| 0 |
365 |
0.384085259396 |
0.633711204665 |
0.331726873354 |
| 1 |
730 |
0.0432828906933 |
0.22382813027 |
0.0267584698433 |
Boosting reports impurity-based variable importance (the total reduction in squared error credited to each covariate, normalized to sum to one):
``` python
# Impurity-based variable importance
gbm.variable_importance()
```
PolarsRows5Columns2
|
|
term
str
|
importance
f64
|
| 0 |
wt.loss |
0.409465130649 |
| 1 |
ph.ecog |
0.212804757602 |
| 2 |
age |
0.205707927573 |
| 3 |
sex |
0.11564191931 |
| 4 |
ph.karno |
0.0563802648662 |
# IPC-weighted ridge regression
Not every alternative to Cox needs a tree. `IPCRidge` fits a penalized linear model on log survival times, but instead of the Cox partial likelihood it uses inverse-probability-of-censoring (IPC) weights to correct for censoring. Only uncensored subjects contribute to the fit, and each is weighted by 1 / G(t-) from the Kaplan-Meier censoring estimate so that subjects whose events occur when censoring is heavy receive greater influence. The result is a simple linear model that does not assume proportional hazards.
``` python
ridge = gw.IPCRidge(alpha=1.0).fit(y, lung[covariates])
ridge
```
IPCRidge (IPC-weighted ridge, alpha=1.0)
coef
(Intercept) 5.315
age -0.01992
sex 0.6572
ph.ecog -0.5318
ph.karno 0.01457
wt.loss 0.001947
n = 213, events = 151 (used for fitting)
The `alpha` parameter controls the ridge penalty. Setting `alpha=0` gives unpenalized IPC-weighted OLS. Larger values shrink the coefficients toward zero, which helps when covariates are correlated or numerous.
Predictions are on the log-time scale by default (`type="lp"`). Pass `type="response"` to get predicted survival times on the original scale.
``` python
ridge.predict(lung[covariates][:5], type="response")
```
array([386.72519998, 491.17815835, 280.69111386, 509.60077987,
109.31469154])
The coefficient table is available via [tidy()](../reference/tidy.md#greenwood.tidy), and model-level summaries via [glance()](../reference/glance.md#greenwood.glance).
``` python
gw.tidy(ridge)
```
PolarsRows6Columns2
|
|
term
str
|
estimate
f64
|
| 0 |
(Intercept) |
5.31469425325 |
| 1 |
age |
-0.0199243788921 |
| 2 |
sex |
0.657164509677 |
| 3 |
ph.ecog |
-0.531839726515 |
| 4 |
ph.karno |
0.0145723160522 |
| 5 |
wt.loss |
0.00194699468604 |
``` python
gw.glance(ridge)
```
PolarsRows1Columns4
|
|
n
i64
|
nevent
i64
|
alpha
f64
|
n_features
i64
|
| 0 |
213 |
151 |
1 |
5 |
Because the linear predictor is on the log-time scale (higher means longer survival), negate it to get a risk score compatible with [concordance_index()](../reference/concordance_index.md#greenwood.concordance_index).
``` python
risk = -ridge.predict(lung[covariates], type="lp")
gw.concordance_index(y, risk)
```
0.6435882352941177
# Visualizing per-subject curves
Because these models predict an individual survival curve for every subject, [plot_predicted_survival()](../reference/plot_predicted_survival.md#greenwood.plot_predicted_survival) draws one step curve per row of new data. It works with any of the machine-learning models (and with [CoxPH](../reference/CoxPH.md#greenwood.CoxPH)), so you can compare how predicted survival varies across individuals.
``` python
# Compare predicted survival for four patients from the forest
gw.plot_predicted_survival(rsf, lung[covariates][:4])
```
Pass `type="cumulative_hazard"` for the cumulative-hazard scale, or `labels=` to name the curves.
# Comparing models
All of these models expose a risk score, so you can judge and compare them with the same censoring-aware metrics as any other model. Here we compare the three ensembles by concordance on the training data (for an honest estimate, evaluate on held-out data or use the forest's out-of-bag score).
``` python
# In-sample concordance for the three ensemble models
scores = {
"random forest": gw.concordance_index(y, rsf.predict(lung[covariates])),
"extra trees": gw.concordance_index(y, ext.predict(lung[covariates])),
"gradient boosting": gw.concordance_index(y, gbm.predict(lung[covariates])),
}
scores
```
{'random forest': 0.8101176470588235,
'extra trees': 0.726764705882353,
'gradient boosting': 0.7446764705882353}
See [Prediction performance](prediction-performance.md) for concordance on held-out data, the Brier score, and time-dependent AUC (the right way to compare these models).