A survival decision tree grown with the log-rank splitting rule.
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
# Load data and build a right-censored response
lung = gw.load_dataset("lung", backend="pandas").dropna(subset=["ph.ecog", "ph.karno"])
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))
cols = ["age", "sex", "ph.ecog", "ph.karno", "wt.loss"]
# Fit a survival tree and score the first five subjects
tree = gw.SurvivalTree(max_depth=3, random_state=0).fit(y, lung[cols])
tree.predict(lung[cols])[:5]
array([95.15345288, 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.
fit(
surv,
covariates,
*,
data=None,
_rng=None,
_event_times=None,
_time=None,
_event=None,
)
Parameters
surv: Surv | None-
A right-censored Surv response (built with Surv.right()).
covariates: Any-
A dataframe, a 2-D array, or a right-hand-side formula string evaluated against data.
data: Any = None-
DataFrame used to evaluate a formula string.
Returns
SurvivalTree-
The fitted estimator (for method chaining), with cached attributes such as
event_times_, n_features_in_, and feature_names_in_.
predict()
Predict a risk score, survival curves, or cumulative-hazard curves.
predict(
newdata=None,
*,
type="risk",
times=None,
format=None,
)
Three prediction types are available:
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.
type="survival": Kaplan-Meier survival probabilities S(t \mid x) at times, one column per subject. Returned as a DataFrame.
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.
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.