Demo for survival analysis (regression).

Demo for survival analysis (regression). using Accelerated Failure Time (AFT) model.

import os

import numpy as np
import pandas as pd
from sklearn.model_selection import ShuffleSplit

import xgboost as xgb

# The Veterans' Administration Lung Cancer Trial
# The Statistical Analysis of Failure Time Data by Kalbfleisch J. and Prentice R (1980)
CURRENT_DIR = os.path.dirname(__file__)
df = pd.read_csv(os.path.join(CURRENT_DIR, "../data/veterans_lung_cancer.csv"))
print("Training data:")
print(df)

# Split features and labels
y_lower_bound = df["Survival_label_lower_bound"]
y_upper_bound = df["Survival_label_upper_bound"]
X = df.drop(["Survival_label_lower_bound", "Survival_label_upper_bound"], axis=1)

# Split data into training and validation sets
rs = ShuffleSplit(n_splits=2, test_size=0.7, random_state=0)
train_index, valid_index = next(rs.split(X))
dtrain = xgb.DMatrix(X.values[train_index, :])
dtrain.set_float_info("label_lower_bound", y_lower_bound[train_index])
dtrain.set_float_info("label_upper_bound", y_upper_bound[train_index])
dvalid = xgb.DMatrix(X.values[valid_index, :])
dvalid.set_float_info("label_lower_bound", y_lower_bound[valid_index])
dvalid.set_float_info("label_upper_bound", y_upper_bound[valid_index])

# Train gradient boosted trees using AFT loss and metric
params = {
    "verbosity": 0,
    "objective": "survival:aft",
    "eval_metric": "aft-nloglik",
    "tree_method": "hist",
    "learning_rate": 0.05,
    "aft_loss_distribution": "normal",
    "aft_loss_distribution_scale": 1.20,
    "max_depth": 6,
    "lambda": 0.01,
    "alpha": 0.02,
}
bst = xgb.train(
    params,
    dtrain,
    num_boost_round=10000,
    evals=[(dtrain, "train"), (dvalid, "valid")],
    early_stopping_rounds=50,
)

# Run prediction on the validation set
df = pd.DataFrame(
    {
        "Label (lower bound)": y_lower_bound[valid_index],
        "Label (upper bound)": y_upper_bound[valid_index],
        "Predicted label": bst.predict(dvalid),
    }
)
print(df)
# Show only data points with right-censored labels
print(df[np.isinf(df["Label (upper bound)"])])

# Save trained model
bst.save_model("aft_model.json")