Random Survival Forests › Visualizing Random Survival Forests

Michael BrenndoerferOctober 29, 202612 min read
Machine Learning from ScratchAll chapters
Part 8 of 15Tree-Based Models
Chapter 60Random Survival Forests › Visualizing Random Survival Forests
Reading section 5 of 10 · Visualizing Random Survival Forests · 12 minSection 5 of 10 · 12 min

Section 5 of 10 · 12 min read

Random Survival Forests › Visualizing Random Survival ForestsSection view · Sections in this chapter5 / 10 ↓
Sections in this chapter10
Reading tools Definitions, links & notes

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Notes & comments

Before fitting anything, we should see how the two ingredients of the model work. We simulate a dataset in which the truth is known, and use it to look at the split search and at the terminal-node estimate.

The data generator is a Weibull model in which the log hazard is a nonlinear function of the covariates. The Weibull generator uses a shape parameter and a rate parameter. We distinguish the generator parameters from the earlier tree-correlation symbol ρ\rho: the Weibull shape parameter is a=1.5a = 1.5 and the Weibull rate parameter is b=0.2b = 0.2, so the conditional survival is S(t∣x)=exp⁡(−b taeη(x))S(t \mid x) = \exp(-b\, t^{a} e^{\eta(x)}). The covariates are five independent standard normal variables, the first two of which are informative and the last three of which are noise. The covariate contribution to the log hazard is

η(x)=1.2 x1+1.2 sin⁡(2x1)+1.2 (x22−1),\eta(x) = 1.2\,x_1 + 1.2\,\sin(2 x_1) + 1.2\,(x_2^2 - 1),

which contains a linear term, a periodic term, and a quadratic term, and therefore a linear model in x1x_1 and x2x_2 cannot represent the truth. Censoring is exponential with rate 0.300.30, which produces roughly half of the observations censored at the generated event-time scale.

The quantities are:

  • x1,x2,…,x5x_1, x_2, \ldots, x_5: the five standard-normal covariates, of which only x1x_1 and x2x_2 enter η\eta
  • η(x)\eta(x): the covariate contribution to the log hazard; the Weibull baseline also depends on time
  • aa: the Weibull shape parameter, 1.51.5 here
  • bb: the Weibull rate parameter, 0.20.2 here
  • S(t∣x)S(t \mid x): the conditional survival function of the event time given the covariates

The simulated censoring times are independent of both the covariates and the latent event times. This gives independent censoring in this simulation.

PythonIn [7]
CollapseExpand
def make_survival_data(n, rng, nonlinear=True):
    """Simulate Weibull survival times with a nonlinear log hazard."""
    X = rng.normal(size=(n, 5))
    if nonlinear:
        eta = (
            1.2 * X[:, 0]
            + 1.2 * np.sin(2.0 * X[:, 0])
            + 1.2 * (X[:, 1] ** 2 - 1.0)
        )
    else:
        eta = 1.0 * X[:, 0] + 0.8 * X[:, 1]
    shape, rate = (
        1.5,
        0.20,
    )  # named a and b in the text to avoid clashing with rho and lambda
    u = rng.uniform(size=n)
    event_time = (-np.log(u) / (rate * np.exp(eta))) ** (1.0 / shape)
    censor_time = rng.exponential(scale=1.0 / 0.30, size=n)
    return (
        X,
        np.minimum(event_time, censor_time),
        event_time <= censor_time,
        eta,
    )


rng = np.random.default_rng(20240517)
X_train, t_train, e_train, eta_train = make_survival_data(500, rng)
X_test, t_test, e_test, eta_test = make_survival_data(400, rng)
ConsoleOut [8]
CollapseExpand
training rows: 500, events: 246, censored: 50.8%
test rows:     400, events: 186, censored: 53.5%
median follow-up (training): 0.96

The event counts, censoring percentages, and median follow-up above are computed from the simulated arrays and can vary with the seed. The requested row counts remain 500 for training and 400 for testing. In this simulated example, a few hundred events allow several leaves under the minimum-events rule, while censoring makes the risk-set bookkeeping visible.

The split search

The first ingredient is the log-rank statistic. The function below implements the formula we derived, with the sorted arrays used to compute the risk sets at each distinct event time. The grouping vector is true for the left child and false for the right child; a negative returned statistic indicates fewer left-child events than expected in the summed comparison, and its absolute value is what the split search maximizes.

PythonIn [9]
CollapseExpand
def logrank_statistic(time, event, group):
    """Standardised log-rank statistic for a two-group split."""
    time = np.asarray(time, dtype=float)
    event = np.asarray(event, dtype=bool)
    group = np.asarray(group, dtype=bool)
    n = time.size

    order = np.argsort(time, kind="mergesort")
    t_sorted = time[order]
    e_sorted = event[order].astype(float)
    g_sorted = group[order].astype(float)

    unique_times, first = np.unique(t_sorted, return_index=True)
    at_risk = n - first
    deaths = np.add.reduceat(e_sorted, first)
    deaths_g = np.add.reduceat(e_sorted * g_sorted, first)

    group_before = np.concatenate(([0.0], np.cumsum(g_sorted)))[first]
    at_risk_g = g_sorted.sum() - group_before
    at_risk_o = at_risk - at_risk_g

    keep = deaths > 0
    d, d_g = deaths[keep], deaths_g[keep]
    y, y_g, y_o = at_risk[keep], at_risk_g[keep], at_risk_o[keep]

    observed = d_g
    expected = y_g * d / y
    variance = np.zeros_like(d)
    ok = y > 1
    variance[ok] = (
        d[ok]
        * (y[ok] - d[ok])
        * y_g[ok]
        * y_o[ok]
        / (y[ok] ** 2 * (y[ok] - 1.0))
    )

    total_variance = variance.sum()
    if total_variance <= 0:
        return 0.0
    return float((observed - expected).sum() / np.sqrt(total_variance))

The keep filter is important: the statistic is a sum over event times, because a time at which nobody dies contributes nothing to either the numerator or the variance. The ok guard handles any event time at which the risk set has only a single person, where the displayed variance fraction is undefined and the comparison variance is zero.

The split search sweeps over candidate thresholds. For each feature we take a grid of quantiles, form the group membership, and evaluate the statistic. This is deliberately the same kind of search a tree performs, restricted here to a coarse grid.

PythonIn [10]
CollapseExpand
def best_split(X, time, event, features, n_cuts=10, min_leaf=15, min_events=8):
    """Search a quantile grid on each candidate feature; return the best log-rank split."""
    best = None
    for feature in features:
        x = X[:, feature]
        for cut in np.unique(np.quantile(x, np.linspace(0.1, 0.9, n_cuts))):
            left = x <= cut
            n_left = int(left.sum())
            if n_left < min_leaf or x.size - n_left < min_leaf:
                continue
            d_left = event[left].sum()
            if d_left < min_events or event.sum() - d_left < min_events:
                continue
            stat = logrank_statistic(time, event, left)
            if best is None or abs(stat) > abs(best[2]):
                best = (int(feature), float(cut), stat)
    return best

The two figures below sweep thresholds on the two informative features and plot the signed log-rank statistic. On this seeded sample, the x1x_1 curve stays negative and reaches its largest magnitude at an interior threshold near 0.530.53. The x2x_2 curve changes sign near the centre, with larger magnitudes toward the tails. Its largest magnitude on this grid occurs near 1.561.56, separating the upper tail from the remaining observations. The split search compares absolute values, not the heights of the signed curves.

PythonIn [11]
CollapseExpand
candidates_x1 = np.quantile(X_train[:, 0], np.linspace(0.02, 0.98, 60))
candidates_x2 = np.quantile(X_train[:, 1], np.linspace(0.02, 0.98, 60))
stats_x1 = np.array(
    [
        logrank_statistic(t_train, e_train, X_train[:, 0] <= c)
        for c in candidates_x1
    ]
)
stats_x2 = np.array(
    [
        logrank_statistic(t_train, e_train, X_train[:, 1] <= c)
        for c in candidates_x2
    ]
)
PythonIn [12]
CollapseExpand
import matplotlib.pyplot as plt
from book_plot_style import PALETTE, polish_axes, theme_color, use_book_style


def pcol(name, fallback=0):
    try:
        return PALETTE[name]
    except (KeyError, TypeError, IndexError):
        return plt.rcParams["axes.prop_cycle"].by_key()["color"][fallback]


use_book_style()

plt.figure(figsize=(6.0, 4.0))
ax = plt.gca()
ax.plot(candidates_x1, stats_x1, color=PALETTE["primary"], linewidth=1.1)
ax.axhline(0.0, color=PALETTE["muted"], linewidth=0.8, linestyle=":")
k = int(np.argmax(np.abs(stats_x1)))
ax.plot(
    candidates_x1[k],
    stats_x1[k],
    "o",
    markerfacecolor=PALETTE["paper"],
    markeredgecolor=PALETTE["primary"],
    markersize=8,
)
ax.axvline(
    candidates_x1[k], color=PALETTE["primary"], linestyle="--", linewidth=1.0
)
ax.set_xlabel("candidate threshold on $x_1$")
ax.set_ylabel("log-rank statistic $L(s)$")
ax.set_title("Split search on the linear-plus-sine feature")
polish_axes(ax)
plt.show()
VisualizationOut [12]
CollapseExpand
A line plot of the signed log-rank statistic against candidate thresholds on x1. The curve stays below zero and has a broad interior trough, and an open marker with a dashed vertical line shows the selected threshold.
Log-rank statistic across every candidate threshold on x1. Each point is one threshold from a quantile grid, and the open marker with a dashed line marks the value a tree would select. The curve stays negative and has a broad interior trough, with the selected threshold near the minimum, showing that the split search compares absolute values rather than signed heights.
PythonIn [13]
CollapseExpand
import matplotlib.pyplot as plt
from book_plot_style import PALETTE, polish_axes, theme_color, use_book_style


def pcol(name, fallback=0):
    try:
        return PALETTE[name]
    except (KeyError, TypeError, IndexError):
        return plt.rcParams["axes.prop_cycle"].by_key()["color"][fallback]


use_book_style()

plt.figure(figsize=(6.0, 4.0))
ax = plt.gca()
ax.plot(candidates_x2, stats_x2, color=PALETTE["amber"], linewidth=1.1)
ax.axhline(0.0, color=PALETTE["muted"], linewidth=0.8, linestyle=":")
k = int(np.argmax(np.abs(stats_x2)))
ax.plot(
    candidates_x2[k],
    stats_x2[k],
    "o",
    markerfacecolor=PALETTE["paper"],
    markeredgecolor=PALETTE["amber"],
    markersize=8,
)
ax.axvline(
    candidates_x2[k], color=PALETTE["amber"], linestyle="--", linewidth=1.0
)
ax.set_xlabel("candidate threshold on $x_2$")
ax.set_ylabel("log-rank statistic $L(s)$")
ax.set_title("Split search on the quadratic feature")
polish_axes(ax)
plt.show()
VisualizationOut [13]
CollapseExpand
A line plot of the signed log-rank statistic against candidate thresholds on x2. The curve moves from positive to negative and reaches its largest absolute value in the upper tail, and an open marker with a dashed vertical line shows the selected threshold.
Log-rank statistic across every candidate threshold on x2. Each point is one threshold from a quantile grid, and the open marker with a dashed line marks the value a tree would select. The curve changes sign near the center and reaches its largest magnitude in the upper tail, where the selected threshold isolates the high-risk tail and leaves the opposite tail in the other child.

The x1x_1 curve has a broad interior trough rather than a monotone shape. The selected threshold near x1=0.53x_1=0.53 separates a lower-risk left child from a higher-risk right child in this sample. The x2x_2 curve shows a threshold near x2=1.56x_2=1.56 that isolates the upper high-risk tail. A single threshold cannot put both tails of a quadratic effect into one child while leaving the centre in the other. Later splits can isolate the two tails in separate leaves, allowing their predictions to differ from the centre. Bootstrap samples and random feature subsets can change the selected thresholds across trees.

What a terminal node stores

The second ingredient is the Nelson–Aalen estimator. The function below sorts the data, groups the times, and accumulates the increments. It returns the distinct event times and the cumulative hazard at each of them, so the result is a step function that can be evaluated anywhere.

PythonIn [14]
CollapseExpand
def nelson_aalen(time, event):
    """Nelson-Aalen estimate of the cumulative hazard: event times and cumulative hazard."""
    time = np.asarray(time, dtype=float)
    event = np.asarray(event, dtype=bool)
    n = time.size
    if n == 0:
        return np.empty(0), np.empty(0)

    order = np.argsort(time, kind="mergesort")
    t_sorted, e_sorted = time[order], event[order]

    unique_times, first = np.unique(t_sorted, return_index=True)
    at_risk = n - first
    deaths = np.add.reduceat(e_sorted.astype(float), first)

    keep = deaths > 0
    return unique_times[keep], np.cumsum(deaths[keep] / at_risk[keep])

The at_risk = n - first line is the whole risk-set computation in disguise: first is the index of the first occurrence of each distinct time in the sorted arrays, so n - first is the number of observations whose time is at least that value. It is the same as counting the patients with Ti≥tT_i \ge t, and it takes linear time after sorting, in addition to the sorting cost.

Now we take the best split found on the full data, form the two children, and estimate the cumulative hazard in each. The visual sweeps above use a finer, wider quantile grid than this helper; here the root search also enforces at least 25 observations and 15 events in each child, so its selected threshold can differ from the sweep markers.

PythonIn [15]
CollapseExpand
root_feature, root_cut, root_stat = best_split(
    X_train,
    t_train,
    e_train,
    features=range(X_train.shape[1]),
    min_leaf=25,
    min_events=15,
)
in_left = X_train[:, root_feature] <= root_cut
times_all, cumhaz_all = nelson_aalen(t_train, e_train)
times_left, cumhaz_left = nelson_aalen(t_train[in_left], e_train[in_left])
times_right, cumhaz_right = nelson_aalen(t_train[~in_left], e_train[~in_left])
plot_times_all = np.r_[0.0, times_all]
plot_cumhaz_all = np.r_[0.0, cumhaz_all]
plot_times_left = np.r_[0.0, times_left]
plot_cumhaz_left = np.r_[0.0, cumhaz_left]
plot_times_right = np.r_[0.0, times_right]
plot_cumhaz_right = np.r_[0.0, cumhaz_right]
PythonIn [16]
CollapseExpand
import matplotlib.pyplot as plt
from book_plot_style import PALETTE, polish_axes, use_book_style

use_book_style()

plt.figure(figsize=(6.0, 4.6))
ax = plt.gca()
ax.step(
    plot_times_all,
    plot_cumhaz_all,
    where="post",
    color=PALETTE["muted"],
    linewidth=1.1,
    label="all patients",
)
ax.step(
    plot_times_left,
    plot_cumhaz_left,
    where="post",
    color=PALETTE["primary"],
    linewidth=1.1,
    label=f"$x_{{{root_feature + 1}}} \\leq {root_cut:.2f}$",
)
ax.step(
    plot_times_right,
    plot_cumhaz_right,
    where="post",
    color=PALETTE["amber"],
    linewidth=1.1,
    label=f"$x_{{{root_feature + 1}}} > {root_cut:.2f}$",
)
ax.set_xlabel("time")
ax.set_ylabel("cumulative hazard $\\hat{H}(t)$")
ax.set_xlim(0.0, 6.0)
ax.set_title("Nelson–Aalen cumulative hazard by child")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.20), ncol=3)
polish_axes(ax)
plt.show()
VisualizationOut [16]
CollapseExpand
Three cumulative hazard step functions: a muted reference curve for all patients, a blue curve for the low-risk child that stays low, and an orange curve for the high-risk child that climbs steeply.
Nelson–Aalen cumulative hazard staircases for the two children of a single split. The muted reference curve is the whole sample, the blue curve is the left child, and the orange curve is the right child. The children separate immediately, showing that each terminal node stores a step function rather than a single number, with jump sizes set by the node's own risk sets.
PythonIn [17]
CollapseExpand
import matplotlib.pyplot as plt
from book_plot_style import PALETTE, polish_axes, use_book_style

use_book_style()

plt.figure(figsize=(6.0, 4.6))
ax = plt.gca()
ax.step(
    plot_times_all,
    np.exp(-plot_cumhaz_all),
    where="post",
    color=PALETTE["muted"],
    linewidth=1.1,
    label="all patients",
)
ax.step(
    plot_times_left,
    np.exp(-plot_cumhaz_left),
    where="post",
    color=PALETTE["primary"],
    linewidth=1.1,
    label=f"$x_{{{root_feature + 1}}} \\leq {root_cut:.2f}$",
)
ax.step(
    plot_times_right,
    np.exp(-plot_cumhaz_right),
    where="post",
    color=PALETTE["amber"],
    linewidth=1.1,
    label=f"$x_{{{root_feature + 1}}} > {root_cut:.2f}$",
)
ax.set_xlabel("time")
ax.set_ylabel("survival $\\hat{S}(t) = \\exp(-\\hat{H}(t))$")
ax.set_xlim(0.0, 6.0)
ax.set_title("Survival is the exponential of the negative")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.20), ncol=3)
polish_axes(ax)
plt.show()
VisualizationOut [17]
CollapseExpand
Three survival step curves that start at one: a blue curve for the low-risk child that declines slowly, an orange curve for the high-risk child that declines quickly, and a muted reference curve for all patients in between.
Implied survival functions for the two children of a single split, obtained by exponentiating the negative of the Nelson–Aalen cumulative hazard. The high-risk child declines faster than the low-risk child, and the whole-sample curve lies between them, reflecting the same partition seen in the cumulative-hazard figure.

The fitted child staircases separate in this sample; a selected split can also separate curves by chance. The jumps are visible: each event time in a node produces a vertical increment of size d/Yd/Y, and their sizes depend on the local event count and number still at risk. A minimum event count does not itself bound the last jump. If a leaf had only one event and one patient, the jump would be a full unit of hazard, and the survival function would collapse to e−1e^{-1} in one step. Our implementation constrains both observations and events per child. Other implementations, including the reference used below, express their leaf-size constraints differently.

Comments

Add to the conversation

A question, a correction, or an example from your own work.

No comments yet. Be the first to share your thoughts!

Newsletter

Stay up to date

Get new articles, book updates, and notes on AI and technology.

Unsubscribe at any time. Privacy policy