Merge pull request #27 from loft-br/fix/preds-using-early-stop

fix:predictions use best tree when early stopping

Author: GabrielGimenez

tests/test_survival_curves.py

@@ -67,6 +67,23 @@ def test_survival_curve(model):
     assert_survival_curve(xgbse, X_test, preds, cindex)
 
 
+@pytest.mark.parametrize(
+    "model", [XGBSEDebiasedBCE, XGBSEKaplanNeighbors, XGBSEStackedWeibull]
+)
+def test_survival_curve_without_early_stopping(model):
+    xgbse = model()
+
+    xgbse.fit(
+        X_train,
+        y_train,
+    )
+
+    preds = xgbse.predict(X_test)
+    cindex = concordance_index(y_test, preds)
+
+    assert_survival_curve(xgbse, X_test, preds, cindex)
+
+
 def test_survival_curve_tree():
     xgbse = XGBSEKaplanTree()
 

xgbse/_base.py

@@ -58,14 +58,18 @@ def get_neighbors(
             index = self.tree
         else:
             index_matrix = xgb.DMatrix(index_data)
-            index_leaves = self.bst.predict(index_matrix, pred_leaf=True)
+            index_leaves = self.bst.predict(
+                index_matrix, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+            )
 
             if len(index_leaves.shape) == 1:
                 index_leaves = index_leaves.reshape(-1, 1)
             index = BallTree(index_leaves, metric="hamming")
 
         query_matrix = xgb.DMatrix(query_data)
-        query_leaves = self.bst.predict(query_matrix, pred_leaf=True)
+        query_leaves = self.bst.predict(
+            query_matrix, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
 
         if len(query_leaves.shape) == 1:
             query_leaves = query_leaves.reshape(-1, 1)

xgbse/_debiased_bce.py

@@ -228,7 +228,9 @@ def fit(
         self.feature_importances_ = self.bst.get_score()
         # predicting and encoding leaves
         self.encoder = OneHotEncoder()
-        leaves = self.bst.predict(dtrain, pred_leaf=True)
+        leaves = self.bst.predict(
+            dtrain, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
         leaves_encoded = self.encoder.fit_transform(leaves)
 
         # convert targets for using with logistic regression
@@ -244,7 +246,9 @@ def fit(
             if index_id is None:
                 index_id = X.index.copy()
 
-            index_leaves = self.bst.predict(dtrain, pred_leaf=True)
+            index_leaves = self.bst.predict(
+                dtrain, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+            )
             self.tree = BallTree(index_leaves, metric="hamming")
 
         self.index_id = index_id
@@ -369,7 +373,9 @@ def predict(self, X, return_interval_probs=False):
         d_matrix = xgb.DMatrix(X)
 
         # getting leaves and extracting neighbors
-        leaves = self.bst.predict(d_matrix, pred_leaf=True)
+        leaves = self.bst.predict(
+            d_matrix, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
         leaves_encoded = self.encoder.transform(leaves)
 
         # predicting from logistic regression artifacts

xgbse/_kaplan_neighbors.py

@@ -175,7 +175,9 @@ def fit(
         self.feature_importances_ = self.bst.get_score()
 
         # creating nearest neighbor index
-        leaves = self.bst.predict(dtrain, pred_leaf=True)
+        leaves = self.bst.predict(
+            dtrain, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
 
         self.tree = BallTree(leaves, metric="hamming", leaf_size=40)
 
@@ -229,7 +231,9 @@ def predict(
         d_matrix = xgb.DMatrix(X)
 
         # getting leaves and extracting neighbors
-        leaves = self.bst.predict(d_matrix, pred_leaf=True)
+        leaves = self.bst.predict(
+            d_matrix, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
 
         if self.radius:
             assert self.radius > 0, "Radius must be positive"
@@ -394,7 +398,9 @@ def fit(
         self.feature_importances_ = self.bst.get_score()
 
         # getting leaves
-        leaves = self.bst.predict(dtrain, pred_leaf=True)
+        leaves = self.bst.predict(
+            dtrain, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
 
         # organizing elements per leaf
         leaf_neighs = (
@@ -462,7 +468,9 @@ def predict(self, X, return_ci=False, return_interval_probs=False):
         d_matrix = xgb.DMatrix(X)
 
         # getting leaves and extracting neighbors
-        leaves = self.bst.predict(d_matrix, pred_leaf=True)
+        leaves = self.bst.predict(
+            d_matrix, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+        )
 
         # searching for kaplan meier curves in leaves
         preds_df = self._train_survival.loc[leaves].reset_index(drop=True)

xgbse/_stacked_weibull.py

@@ -171,7 +171,7 @@ def fit(
         self.feature_importances_ = self.bst.get_score()
 
         # predicting risk from XGBoost
-        train_risk = self.bst.predict(dtrain)
+        train_risk = self.bst.predict(dtrain, ntree_limit=self.bst.best_ntree_limit)
 
         # replacing 0 by minimum positive value in df
         # so Weibull can be fitted
@@ -192,7 +192,9 @@ def fit(
             if index_id is None:
                 index_id = X.index.copy()
 
-            index_leaves = self.bst.predict(dtrain, pred_leaf=True)
+            index_leaves = self.bst.predict(
+                dtrain, pred_leaf=True, ntree_limit=self.bst.best_ntree_limit
+            )
             self.tree = BallTree(index_leaves, metric="hamming")
 
         self.index_id = index_id
@@ -222,7 +224,7 @@ def predict(self, X, return_interval_probs=False):
         d_matrix = xgb.DMatrix(X)
 
         # getting leaves and extracting neighbors
-        risk = self.bst.predict(d_matrix)
+        risk = self.bst.predict(d_matrix, ntree_limit=self.bst.best_ntree_limit)
         weibull_score_df = pd.DataFrame({"risk": risk})
 
         # predicting from logistic regression artifacts