Add Accelerated Failure Time loss for survival analysis task

R-package/R/xgb.DMatrix.R

@@ -243,6 +243,18 @@ setinfo.xgb.DMatrix <- function(object, name, info, ...) {
     .Call(XGDMatrixSetInfo_R, object, name, as.numeric(info))
     return(TRUE)
   }
+  if (name == "label_lower_bound") {
+    if (length(info) != nrow(object))
+      stop("The length of lower-bound labels must equal to the number of rows in the input data")
+    .Call(XGDMatrixSetInfo_R, object, name, as.numeric(info))
+    return(TRUE)
+  }
+  if (name == "label_upper_bound") {
+    if (length(info) != nrow(object))
+      stop("The length of upper-bound labels must equal to the number of rows in the input data")
+    .Call(XGDMatrixSetInfo_R, object, name, as.numeric(info))
+    return(TRUE)
+  }
   if (name == "weight") {
     if (length(info) != nrow(object))
       stop("The length of weights must equal to the number of rows in the input data")

amalgamation/xgboost-all0.cc

@@ -14,13 +14,15 @@
 #include "../src/metric/elementwise_metric.cc"
 #include "../src/metric/multiclass_metric.cc"
 #include "../src/metric/rank_metric.cc"
+#include "../src/metric/survival_metric.cc"
 
 // objectives
 #include "../src/objective/objective.cc"
 #include "../src/objective/regression_obj.cc"
 #include "../src/objective/multiclass_obj.cc"
 #include "../src/objective/rank_obj.cc"
 #include "../src/objective/hinge.cc"
+#include "../src/objective/aft_obj.cc"
 
 // gbms
 #include "../src/gbm/gbm.cc"
@@ -44,7 +46,7 @@
 #include "../src/data/sparse_page_dmatrix.cc"
 #endif
 
-// tress
+// trees
 #include "../src/tree/param.cc"
 #include "../src/tree/split_evaluator.cc"
 #include "../src/tree/tree_model.cc"
@@ -72,6 +74,8 @@
 #include "../src/common/hist_util.cc"
 #include "../src/common/json.cc"
 #include "../src/common/io.cc"
+#include "../src/common/survival_util.cc"
+#include "../src/common/probability_distribution.cc"
 #include "../src/common/version.cc"
 
 // c_api

demo/aft_survival/aft_survival_demo.py

@@ -0,0 +1,54 @@
+"""
+Demo for survival analysis (regression) using Accelerated Failure Time (AFT) model
+"""
+from sklearn.model_selection import ShuffleSplit
+import pandas as pd
+import numpy as np
+import xgboost as xgb
+
+# The Veterans' Administration Lung Cancer Trial
+# The Statistical Analysis of Failure Time Data by Kalbfleisch J. and Prentice R (1980)
+df = pd.read_csv('../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=.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')

demo/aft_survival/aft_survival_demo_with_optuna.py

@@ -0,0 +1,78 @@
+"""
+Demo for survival analysis (regression) using Accelerated Failure Time (AFT) model, using Optuna
+to tune hyperparameters
+"""
+from sklearn.model_selection import ShuffleSplit
+import pandas as pd
+import numpy as np
+import xgboost as xgb
+import optuna
+
+# The Veterans' Administration Lung Cancer Trial
+# The Statistical Analysis of Failure Time Data by Kalbfleisch J. and Prentice R (1980)
+df = pd.read_csv('../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=.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])
+
+# Define hyperparameter search space
+base_params = {'verbosity': 0,
+              'objective': 'survival:aft',
+              'eval_metric': 'aft-nloglik',
+              'tree_method': 'hist'}  # Hyperparameters common to all trials
+def objective(trial):
+    params = {'learning_rate': trial.suggest_loguniform('learning_rate', 0.01, 1.0),
+              'aft_loss_distribution': trial.suggest_categorical('aft_loss_distribution',
+                                                                  ['normal', 'logistic', 'extreme']),
+              'aft_loss_distribution_scale': trial.suggest_loguniform('aft_loss_distribution_scale', 0.1, 10.0),
+              'max_depth': trial.suggest_int('max_depth', 3, 8),
+              'lambda': trial.suggest_loguniform('lambda', 1e-8, 1.0),
+              'alpha': trial.suggest_loguniform('alpha', 1e-8, 1.0)}  # Search space
+    params.update(base_params)
+    pruning_callback = optuna.integration.XGBoostPruningCallback(trial, 'valid-aft-nloglik')
+    bst = xgb.train(params, dtrain, num_boost_round=10000,
+                    evals=[(dtrain, 'train'), (dvalid, 'valid')], 
+                    early_stopping_rounds=50, verbose_eval=False, callbacks=[pruning_callback])
+    if bst.best_iteration >= 25:
+        return bst.best_score
+    else:
+        return np.inf  # Reject models with < 25 trees
+
+# Run hyperparameter search
+study = optuna.create_study(direction='minimize')
+study.optimize(objective, n_trials=200)
+print('Completed hyperparameter tuning with best aft-nloglik = {}.'.format(study.best_trial.value))
+params = {}
+params.update(base_params)
+params.update(study.best_trial.params)
+
+# Re-run training with the best hyperparameter combination
+print('Re-running the best trial... params = {}'.format(params))
+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_best_model.json')

demo/aft_survival/aft_survival_viz_demo.py

@@ -0,0 +1,97 @@
+"""
+Visual demo for survival analysis (regression) with Accelerated Failure Time (AFT) model.
+
+This demo uses 1D toy data and visualizes how XGBoost fits a tree ensemble. The ensemble model
+starts out as a flat line and evolves into a step function in order to account for all ranged
+labels.
+"""
+import numpy as np
+import xgboost as xgb
+import matplotlib.pyplot as plt
+
+plt.rcParams.update({'font.size': 13})
+
+# Function to visualize censored labels
+def plot_censored_labels(X, y_lower, y_upper):
+    def replace_inf(x, target_value):
+        x[np.isinf(x)] = target_value
+        return x
+    plt.plot(X, y_lower, 'o', label='y_lower', color='blue')
+    plt.plot(X, y_upper, 'o', label='y_upper', color='fuchsia')
+    plt.vlines(X, ymin=replace_inf(y_lower, 0.01), ymax=replace_inf(y_upper, 1000),
+               label='Range for y', color='gray')
+
+# Toy data
+X = np.array([1, 2, 3, 4, 5]).reshape((-1, 1))
+INF = np.inf
+y_lower = np.array([ 10,  15, -INF, 30, 100])
+y_upper = np.array([INF, INF,   20, 50, INF])
+
+# Visualize toy data
+plt.figure(figsize=(5, 4))
+plot_censored_labels(X, y_lower, y_upper)
+plt.ylim((6, 200))
+plt.legend(loc='lower right')
+plt.title('Toy data')
+plt.xlabel('Input feature')
+plt.ylabel('Label')
+plt.yscale('log')
+plt.tight_layout()
+plt.show(block=True)
+
+# Will be used to visualize XGBoost model
+grid_pts = np.linspace(0.8, 5.2, 1000).reshape((-1, 1))
+
+# Train AFT model using XGBoost
+dmat = xgb.DMatrix(X)
+dmat.set_float_info('label_lower_bound', y_lower)
+dmat.set_float_info('label_upper_bound', y_upper)
+params = {'max_depth': 3, 'objective':'survival:aft', 'min_child_weight': 0}
+
+accuracy_history = []
+def plot_intermediate_model_callback(env):
+    """Custom callback to plot intermediate models"""
+    # Compute y_pred = prediction using the intermediate model, at current boosting iteration
+    y_pred = env.model.predict(dmat)
+    # "Accuracy" = the number of data points whose ranged label (y_lower, y_upper) includes
+    #              the corresponding predicted label (y_pred)
+    acc = np.sum(np.logical_and(y_pred >= y_lower, y_pred <= y_upper)/len(X) * 100)
+    accuracy_history.append(acc)
+
+    # Plot ranged labels as well as predictions by the model
+    plt.subplot(5, 3, env.iteration + 1)
+    plot_censored_labels(X, y_lower, y_upper)
+    y_pred_grid_pts = env.model.predict(xgb.DMatrix(grid_pts))
+    plt.plot(grid_pts, y_pred_grid_pts, 'r-', label='XGBoost AFT model', linewidth=4)
+    plt.title('Iteration {}'.format(env.iteration), x=0.5, y=0.8)
+    plt.xlim((0.8, 5.2))
+    plt.ylim((1 if np.min(y_pred) < 6 else 6, 200))
+    plt.yscale('log')
+
+res = {}
+plt.figure(figsize=(12,13))
+bst = xgb.train(params, dmat, 15, [(dmat, 'train')], evals_result=res,
+                callbacks=[plot_intermediate_model_callback])
+plt.tight_layout()
+plt.legend(loc='lower center', ncol=4,
+           bbox_to_anchor=(0.5, 0),
+           bbox_transform=plt.gcf().transFigure)
+plt.tight_layout()
+
+# Plot negative log likelihood over boosting iterations
+plt.figure(figsize=(8,3))
+plt.subplot(1, 2, 1)
+plt.plot(res['train']['aft-nloglik'], 'b-o', label='aft-nloglik')
+plt.xlabel('# Boosting Iterations')
+plt.legend(loc='best')
+
+# Plot "accuracy" over boosting iterations
+# "Accuracy" = the number of data points whose ranged label (y_lower, y_upper) includes
+#              the corresponding predicted label (y_pred)
+plt.subplot(1, 2, 2)
+plt.plot(accuracy_history, 'r-o', label='Accuracy (%)')
+plt.xlabel('# Boosting Iterations')
+plt.legend(loc='best')
+plt.tight_layout()
+
+plt.show()