handle NULL objective in xgb_predict() (#875)

Author: simonpcouch

NEWS.md

@@ -2,6 +2,8 @@
 
 * `.organize_glmnet_pred()` now expects predictions for a single penalty value (#876).
 
+* Fixed bug with prediction from a boosted tree model fitted with `"xgboost"` using a custom objective function (#875).
+
 # parsnip 1.0.4
 
 * For censored regression models, a "reverse Kaplan-Meier" curve is computed for the censoring distribution. This can be used when evaluating this type of model (#855).

R/boost_tree.R

@@ -399,7 +399,7 @@ xgb_predict <- function(object, new_data, ...) {
   res <- predict(object, new_data, ...)
 
   x <- switch(
-    object$params$objective,
+    object$params$objective %||% 3L,
     "binary:logitraw" = stats::binomial()$linkinv(res),
     "multi:softprob" = matrix(res, ncol = object$params$num_class, byrow = TRUE),
     res)

tests/testthat/test_boost_tree_xgboost.R

@@ -217,8 +217,27 @@ test_that('xgboost alternate objective', {
 
   xgb_fit <- spec %>% fit(mpg ~ ., data = mtcars)
   expect_equal(extract_fit_engine(xgb_fit)$params$objective, "reg:pseudohubererror")
-})
+  expect_no_error(xgb_preds <- predict(xgb_fit, new_data = mtcars[1,]))
+  expect_s3_class(xgb_preds, "data.frame")
+
+  logregobj <- function(preds, dtrain) {
+    labels <- xgboost::getinfo(dtrain, "label")
+    preds <- 1 / (1 + exp(-preds))
+    grad <- preds - labels
+    hess <- preds * (1 - preds)
+    return(list(grad = grad, hess = hess))
+  }
+
+  spec2 <-
+    boost_tree() %>%
+    set_engine("xgboost", objective = logregobj) %>%
+    set_mode("classification")
 
+  xgb_fit2 <- spec2 %>% fit(vs ~ ., data = mtcars %>% mutate(vs = as.factor(vs)))
+  expect_equal(rlang::eval_tidy(xgb_fit2$spec$eng_args$objective), logregobj)
+  expect_no_error(xgb_preds2 <- predict(xgb_fit2, new_data = mtcars[1,-8]))
+  expect_s3_class(xgb_preds2, "data.frame")
+})
 
 test_that('submodel prediction', {