respect best_iteration for early-stopped xgboost models

NEWS.md

@@ -1,5 +1,10 @@
 # glex 0.6.0.9000 (development version)
 
+* `glex()` on `xgboost` models fit with early stopping now decomposes only the trees
+  up to `best_iteration`, matching what `predict()` evaluates by default. Previously
+  all fitted trees were decomposed, so the components did not sum to the prediction.
+  Closes #42.
+
 * `randomPlantedForest (>= 0.3.0)` is now required (in `Suggests:`): it fixes an
   out-of-bounds read in `purify_3()` that crashed R on Windows
   (PlantedML/randomPlantedForest#61), so rpf tests and examples run on all platforms.

R/glex.R

@@ -217,6 +217,19 @@ glex.xgb.Booster <- function(
   trees <- xgboost::xgb.model.dt.tree(model = object, use_int_id = TRUE)
   trees$Type <- "<"
 
+  # Early stopping stores the 0-based best round as a booster attribute and
+  # predict() defaults to using only trees up to it; decompose the same model
+  # predict() evaluates. A round is several trees for multiclass models and
+  # num_parallel_tree > 1, so translate rounds to trees via their ratio.
+  best_iteration <- xgboost::xgb.attributes(object)$best_iteration
+  if (!is.null(best_iteration)) {
+    n_rounds <- xgboost::xgb.get.num.boosted.rounds(object)
+    trees_per_round <- length(unique(trees$Tree)) / n_rounds
+    trees <- trees[
+      trees$Tree < (as.integer(best_iteration) + 1) * trees_per_round,
+    ]
+  }
+
   # Calculate components
   res <- calc_components(
     trees,

tests/testthat/test-glex-xgboost.R

@@ -330,3 +330,41 @@ test_that("xgboost reg:gamma (log link) is reconstructed on margin and response
   expect_equal(response_from_shap, unname(pred_response), tolerance = 1e-5)
   expect_equal(response_from_m, unname(pred_response), tolerance = 1e-5)
 })
+
+test_that("early-stopped models are decomposed up to best_iteration, like predict()", {
+  set.seed(1)
+  n <- 200
+  x <- matrix(rnorm(n * 4), ncol = 4, dimnames = list(NULL, paste0("x", 1:4)))
+  y <- x[, 1] + rnorm(n, sd = 3)
+  dtrain <- xgboost::xgb.DMatrix(x[1:140, ], label = y[1:140], nthread = 1)
+  deval <- xgboost::xgb.DMatrix(x[141:200, ], label = y[141:200], nthread = 1)
+
+  bst <- xgboost::xgb.train(
+    params = xgboost::xgb.params(max_depth = 3, learning_rate = 0.5, nthread = 1),
+    data = dtrain,
+    nrounds = 500,
+    evals = list(eval = deval),
+    early_stopping_rounds = 3,
+    verbose = 0
+  )
+
+  # Premise: early stopping engaged and predict() defaults to best_iteration,
+  # which differs from the full model
+  best <- as.integer(xgboost::xgb.attributes(bst)$best_iteration)
+  expect_lt(best + 1, xgboost::xgb.get.num.boosted.rounds(bst))
+  p_default <- predict(bst, x, outputmargin = TRUE)
+  p_all <- predict(bst, x, outputmargin = TRUE, iterationrange = "all")
+  expect_false(isTRUE(all.equal(p_default, p_all)))
+
+  gl <- glex(bst, x)
+  expect_equal(
+    unname(gl$intercept + rowSums(gl$m)),
+    unname(p_default),
+    tolerance = 1e-5
+  )
+  expect_equal(
+    unname(gl$intercept + rowSums(gl$shap)),
+    unname(p_default),
+    tolerance = 1e-5
+  )
+})