dcf439932a
* [WIP] Add lower and upper bounds on the label for survival analysis * Update test MetaInfo.SaveLoadBinary to account for extra two fields * Don't clear qids_ for version 2 of MetaInfo * Add SetInfo() and GetInfo() method for lower and upper bounds * changes to aft * Add parameter class for AFT; use enum's to represent distribution and event type * Add AFT metric * changes to neg grad to grad * changes to binomial loss * changes to overflow * changes to eps * changes to code refactoring * changes to code refactoring * changes to code refactoring * Re-factor survival analysis * Remove aft namespace * Move function bodies out of AFTNormal and AFTLogistic, to reduce clutter * Move function bodies out of AFTLoss, to reduce clutter * Use smart pointer to store AFTDistribution and AFTLoss * Rename AFTNoiseDistribution enum to AFTDistributionType for clarity The enum class was not a distribution itself but a distribution type * Add AFTDistribution::Create() method for convenience * changes to extreme distribution * changes to extreme distribution * changes to extreme * changes to extreme distribution * changes to left censored * deleted cout * changes to x,mu and sd and code refactoring * changes to print * changes to hessian formula in censored and uncensored * changes to variable names and pow * changes to Logistic Pdf * changes to parameter * Expose lower and upper bound labels to R package * Use example weights; normalize log likelihood metric * changes to CHECK * changes to logistic hessian to standard formula * changes to logistic formula * Comply with coding style guideline * Revert back Rabit submodule * Revert dmlc-core submodule * Comply with coding style guideline (clang-tidy) * Fix an error in AFTLoss::Gradient() * Add missing files to amalgamation * Address @RAMitchell's comment: minimize future change in MetaInfo interface * Fix lint * Fix compilation error on 32-bit target, when size_t == bst_uint * Allocate sufficient memory to hold extra label info * Use OpenMP to speed up * Fix compilation on Windows * Address reviewer's feedback * Add unit tests for probability distributions * Make Metric subclass of Configurable * Address reviewer's feedback: Configure() AFT metric * Add a dummy test for AFT metric configuration * Complete AFT configuration test; remove debugging print * Rename AFT parameters * Clarify test comment * Add a dummy test for AFT loss for uncensored case * Fix a bug in AFT loss for uncensored labels * Complete unit test for AFT loss metric * Simplify unit tests for AFT metric * Add unit test to verify aggregate output from AFT metric * Use EXPECT_* instead of ASSERT_*, so that we run all unit tests * Use aft_loss_param when serializing AFTObj This is to be consistent with AFT metric * Add unit tests for AFT Objective * Fix OpenMP bug; clarify semantics for shared variables used in OpenMP loops * Add comments * Remove AFT prefix from probability distribution; put probability distribution in separate source file * Add comments * Define kPI and kEulerMascheroni in probability_distribution.h * Add probability_distribution.cc to amalgamation * Remove unnecessary diff * Address reviewer's feedback: define variables where they're used * Eliminate all INFs and NANs from AFT loss and gradient * Add demo * Add tutorial * Fix lint * Use 'survival:aft' to be consistent with 'survival:cox' * Move sample data to demo/data * Add visual demo with 1D toy data * Add Python tests Co-authored-by: Philip Cho <chohyu01@cs.washington.edu>
54 lines
2.1 KiB
Python
54 lines
2.1 KiB
Python
"""
|
|
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') |