Quick Start Guide: Optimal Survival Trees

This is a Python version of the corresponding OptimalTrees quick start guide.

In this example we will use Optimal Survival Trees (OST) on the Monoclonal Gammopathy dataset to predict patient survival time (refer to the data preparation guide to learn more about the data format for survival problems). First we load in the data and split into training and test datasets:

import pandas as pd
df = pd.read_csv("mgus2.csv")
df.sex = df.sex.astype('category')

      Unnamed: 0    id  age sex   hgb  ...  mspike  ptime  pstat  futime  death
0              1     1   88   F  13.1  ...     0.5     30      0      30      1
1              2     2   78   F  11.5  ...     2.0     25      0      25      1
2              3     3   94   M  10.5  ...     2.6     46      0      46      1
3              4     4   68   M  15.2  ...     1.2     92      0      92      1
4              5     5   90   F  10.7  ...     1.0      8      0       8      1
5              6     6   90   M  12.9  ...     0.5      4      0       4      1
6              7     7   89   F  10.5  ...     1.3    151      0     151      1
...          ...   ...  ...  ..   ...  ...     ...    ...    ...     ...    ...
1377        1378  1378   56   M  16.1  ...     0.5     59      0      59      0
1378        1379  1379   73   M  15.6  ...     0.5     48      0      48      0
1379        1380  1380   69   M  15.0  ...     0.0     22      0      22      1
1380        1381  1381   78   M  14.1  ...     1.9     35      0      35      0
1381        1382  1382   66   M  12.1  ...     0.0     31      0      31      1
1382        1383  1383   82   F  11.5  ...     2.3     38      1      61      0
1383        1384  1384   79   M   9.6  ...     1.7      6      0       6      1

[1384 rows x 11 columns]
from interpretableai import iai
X = df.iloc[:, 2:-4]
died = df.death == 1
times = df.futime
(train_X, train_died, train_times), (test_X, test_died, test_times) = (
    iai.split_data('survival', X, died, times, seed=12345))

Optimal Survival Trees

We will use a GridSearch to fit an OptimalTreeSurvivalLearner:

grid = iai.GridSearch(
    iai.OptimalTreeSurvivalLearner(
        random_seed=1,
        missingdatamode='separate_class',
        minbucket=15,
    ),
    max_depth=range(1, 3),
)
grid.fit(train_X, train_died, train_times,
         validation_criterion='harrell_c_statistic')
grid.get_learner()
Optimal Trees Visualization

The survival tree shows the Kaplan-Meier survival curve for the population in each node as a solid red line.

In each split node:

  • the dotted green line shows the survival curve of the population in the lower/left child
  • the dotted blue line shows the survival curve of the population in the upper/right child

This means that the distance between the green and blue lines gives an indication on how well the split separates the two groups.

In each leaf node, the dotted black line shows the survival curve of the entire population, which allows us to easily see how the survival outlook for this subpopulation differs from the entire population.

In this example, age and hgb partition the population into three subgroups with distinct survival patterns:

  • Node 3: when age < 67.5 and hgb < 12.25, the population has a survival similar to the overall average
  • Node 4: when age < 67.5 and hgb > 12.25, the survival is significantly better than average
  • Node 6: when 67.5 < age < 77.5, the survival is similar to the overall average
  • Node 7: when age > 77.5, the survival is significantly worse than average

We can make predictions on new data using predict. For survival trees, this returns the appropriate SurvivalCurve for each point, which we can then use to query the survival probability for any given time t (in this case for t = 10):

pred_curves = grid.predict(test_X)
t = 10
import numpy as np
np.array([c[t] for c in pred_curves])

array([0.16553847, 0.16553847, 0.09713376, ..., 0.16553847, 0.16553847,
       0.24537037])

Alternatively, you can get this mortality probability for any given time t by providing t as a keyword argument directly:

grid.predict(test_X, t=10)

array([0.16553847, 0.16553847, 0.09713376, ..., 0.16553847, 0.16553847,
       0.24537037])

We can evaluate the quality of the tree using score with any of the supported loss functions. For example, the Harrell's c-statistic on the training set:

grid.score(train_X, train_died, train_times, criterion='harrell_c_statistic')

0.6608294740324181

Or on the test set:

grid.score(test_X, test_died, test_times, criterion='harrell_c_statistic')

0.6668716957368598

We can also evaluate the performance of the tree at a particular point in time using classification criteria. For instance, we can evaluate the AUC of the 10-month survival predictions on the test set:

grid.score(test_X, test_died, test_times, criterion='auc', evaluation_time=10)

0.6793093093093092