SurvivalTree

A survival decision tree grown with the log-rank splitting rule.

Usage

Source

SurvivalTree(
    *,
    max_depth=None,
    min_samples_split=6,
    min_samples_leaf=3,
    max_features=None,
    splitter="best",
    engine="numpy",
    random_state=None,
)

The tree recursively partitions subjects by covariate thresholds, at each step choosing the split that maximizes the two-sample log-rank statistic between the resulting child groups. Each leaf holds a Nelson-Aalen cumulative hazard and a Kaplan-Meier survival curve estimated from its training subjects, evaluated on the full training set’s event-time grid. Prediction routes a subject to a leaf and reads off that leaf’s curves; the scalar risk score is the summed cumulative hazard (ensemble mortality).

A single tree is a high-variance estimator and is most useful as an interpretable building block or the base learner of a RandomSurvivalForest.

Parameters

max_depth: int | None = None

Maximum depth of the tree. None grows until other stopping rules apply.

min_samples_split: int = 6

Minimum number of subjects a node must have to be eligible for splitting.

min_samples_leaf: int = 3

Minimum number of subjects required in each child of a split.

max_features: Any = None

Number of covariates considered at each split: "sqrt", "log2", an int, a float fraction, or None (all features). A random subset is drawn per node.

splitter: str = "best"

How thresholds are chosen: "best" (default) scans every candidate cut-point for the optimal log-rank split; "random" draws a single random threshold per candidate feature (the extremely-randomized-trees rule used by ExtraSurvivalTrees).

engine: str = "numpy"

Compute backend for the best-split search: "numpy" (default) uses the deterministic vectorized path; "numba" uses a compiled kernel (requires the fast extra) that is faster on large datasets; "auto" uses Numba when available. The Numba path produces a statistically equivalent fit but, because it accumulates the log-rank statistic in a different order, may differ from the NumPy path in tie-breaking.

random_state: Any = None

Seed or numpy.random.Generator controlling the per-node feature subsampling.

Examples

Grow a survival tree on the bundled lung dataset and predict a risk score:

import greenwood as gw

lung = gw.load_dataset("lung", backend="pandas").dropna(subset=["ph.ecog", "ph.karno"])

# Fit a survival tree and score the first five subjects
tree = gw.SurvivalTree(max_depth=3, random_state=0).fit(
    "Surv(time, status) ~ age + sex + ph.ecog + ph.karno + wt.loss", data=lung
)
tree.predict(lung)[:5]

array([34.84166667, 95.15345288, 95.15345288, 95.15345288, 95.15345288])

Methods

Name	Description
fit()	Fit the survival tree to a right-censored response and covariate design.
predict()	Predict a risk score, survival curves, or cumulative-hazard curves.

---

fit()

Fit the survival tree to a right-censored response and covariate design.

Usage

Source

fit(
    surv,
    covariates=None,
    *,
    data=None,
    _rng=None,
    _event_times=None,
    _time=None,
    _event=None,
)

Parameters

surv: Surv | Outcome | str | None

A right-censored Surv response, built with gw.Surv(time=..., event=...). An Outcome or a formula string such as 'Surv(time, status == 2) ~ age + sex' is also accepted, with its columns read from data. The right-hand side sets covariates.

covariates: Any = None

A dataframe, a 2-D array, or a right-hand-side formula string evaluated against data. A list of column names in data also works.

data: Any = None

A data frame (pandas, Polars, PyArrow, DuckDB, a lazy frame, …) holding the columns named by the response and covariates. When surv is an Outcome or a formula, rows with a missing value in any column used are dropped before fitting. Covariate formula strings and lists of column names are also resolved here.

Returns

SurvivalTree

The fitted estimator (for method chaining), with cached attributes such as event_times_, n_features_in_, and feature_names_in_.

Examples

Grow a survival tree on the lung dataset with a formula interface. The formula specifies both the response and the covariates in a single string:

import greenwood as gw

lung = gw.load_dataset("lung", backend="pandas").dropna(
    subset=["ph.ecog", "ph.karno"]
)

tree = gw.SurvivalTree(max_depth=3, random_state=0).fit(
    "Surv(time, status == 2) ~ age + sex + ph.ecog + ph.karno",
    data=lung,
)
tree

SurvivalTree (log-rank splits, 15 nodes, 8 leaves)
n = 226, events = 163, features = 4

The same fit can be done by passing an Outcome and the covariate columns separately:

death = gw.Outcome.surv(time="time", event="status == 2")
tree2 = gw.SurvivalTree(max_depth=3, random_state=0).fit(
    death, covariates=["age", "sex", "ph.ecog", "ph.karno"], data=lung
)
tree2

SurvivalTree (log-rank splits, 15 nodes, 8 leaves)
n = 226, events = 163, features = 4

---

predict()

Predict a risk score, survival curves, or cumulative-hazard curves.

Usage

Source

predict(
    newdata=None,
    *,
    type="risk",
    times=None,
    format=None,
)

Three prediction types are available:

1. type="risk" (default): Ishwaran’s ensemble mortality, the sum of the leaf cumulative hazard over the training event times. Higher values indicate higher risk. Returned as a 1-D NumPy array.

2. type="survival": Kaplan-Meier survival probabilities S(t \mid x) at times, one column per subject. Returned as a DataFrame.

3. type="cumulative_hazard": Nelson-Aalen cumulative hazard H(t \mid x) at times. Returned as a DataFrame.

Parameters

newdata: Any = None

Covariates to predict for. None predicts for the training subjects. A data frame may also hold other columns. The covariates are picked out by name and coded as they were at fit time, so the frame the model was fit on can be passed as is.

type: str = "risk"

One of "risk", "survival", or "cumulative_hazard".

times: Any = None

Times at which to evaluate survival / cumulative hazard. Defaults to the training event times. Ignored for type="risk".

format: str | None = None

DataFrame backend for curve output: None, "pandas", "polars", or "pyarrow".

Returns

numpy.ndarray or DataFrame

A risk vector, or a curve frame with a time column and one column per subject.

Examples

Compute the ensemble-mortality risk score for each subject (a 1-D array where higher values indicate higher risk):

import greenwood as gw

lung = gw.load_dataset("lung", backend="pandas").dropna(
    subset=["ph.ecog", "ph.karno"]
)

tree = gw.SurvivalTree(max_depth=3, random_state=0).fit(
    "Surv(time, status) ~ age + sex + ph.ecog + ph.karno + wt.loss", data=lung
)
tree.predict(lung)[:5]

array([34.84166667, 95.15345288, 95.15345288, 95.15345288, 95.15345288])

Predict survival curves at specific time points. The result is a DataFrame with a time column and one column per subject:

tree.predict(lung, type="survival", times=[180, 365, 730])

shape: (3, 227)

time	subject_1	subject_2	subject_3	subject_4	subject_5	subject_6	subject_7	subject_8	subject_9	subject_10	subject_11	subject_12	subject_13	subject_14	subject_15	subject_16	subject_17	subject_18	subject_19	subject_20	subject_21	subject_22	subject_23	subject_24	subject_25	subject_26	subject_27	subject_28	subject_29	subject_30	subject_31	subject_32	subject_33	subject_34	subject_35	subject_36	…	subject_190	subject_191	subject_192	subject_193	subject_194	subject_195	subject_196	subject_197	subject_198	subject_199	subject_200	subject_201	subject_202	subject_203	subject_204	subject_205	subject_206	subject_207	subject_208	subject_209	subject_210	subject_211	subject_212	subject_213	subject_214	subject_215	subject_216	subject_217	subject_218	subject_219	subject_220	subject_221	subject_222	subject_223	subject_224	subject_225	subject_226
f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	…	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64	f64
180.0	0.888889	0.765526	0.765526	0.765526	0.765526	0.765526	0.722222	0.722222	0.765526	0.333333	0.765526	0.888889	0.891892	0.765526	0.765526	0.765526	0.888889	0.722222	0.888889	0.765526	0.866667	0.765526	0.765526	0.765526	0.891892	0.765526	0.277778	0.765526	0.277778	0.866667	0.277778	0.277778	0.722222	0.277778	0.888889	0.888889	…	0.765526	0.888889	0.888889	0.765526	0.765526	0.765526	0.866667	0.765526	0.866667	0.765526	0.891892	0.765526	0.866667	0.866667	0.866667	0.891892	0.722222	0.888889	0.891892	0.866667	0.765526	0.765526	0.888889	0.765526	0.765526	0.866667	0.277778	0.722222	0.891892	0.765526	0.866667	0.5	0.5	0.765526	0.722222	0.765526	0.866667
365.0	0.777778	0.425739	0.425739	0.425739	0.425739	0.425739	0.084259	0.084259	0.425739	0.0	0.425739	0.777778	0.680375	0.425739	0.425739	0.425739	0.777778	0.084259	0.777778	0.425739	0.536022	0.425739	0.425739	0.425739	0.680375	0.425739	0.055556	0.425739	0.055556	0.536022	0.055556	0.055556	0.084259	0.055556	0.777778	0.777778	…	0.425739	0.777778	0.777778	0.425739	0.425739	0.425739	0.536022	0.425739	0.536022	0.425739	0.680375	0.425739	0.536022	0.536022	0.536022	0.680375	0.084259	0.777778	0.680375	0.536022	0.425739	0.425739	0.777778	0.425739	0.425739	0.536022	0.055556	0.084259	0.680375	0.425739	0.536022	0.0	0.0	0.425739	0.084259	0.425739	0.536022
730.0	0.207407	0.101691	0.101691	0.101691	0.101691	0.101691	0.0	0.0	0.101691	0.0	0.101691	0.207407	0.355581	0.101691	0.101691	0.101691	0.207407	0.0	0.207407	0.101691	0.0	0.101691	0.101691	0.101691	0.355581	0.101691	0.0	0.101691	0.0	0.0	0.0	0.0	0.0	0.0	0.207407	0.207407	…	0.101691	0.207407	0.207407	0.101691	0.101691	0.101691	0.0	0.101691	0.0	0.101691	0.355581	0.101691	0.0	0.0	0.0	0.355581	0.0	0.207407	0.355581	0.0	0.101691	0.101691	0.207407	0.101691	0.101691	0.0	0.0	0.0	0.355581	0.101691	0.0	0.0	0.0	0.101691	0.0	0.101691	0.0