allow 'objective' for xgboost to be passed as an engine argument

Author: topepo

NEWS.md

@@ -1,5 +1,7 @@
 # parsnip (development version)
 
+ * For `xgboost` models, users can now pass `objective` to `set_engine("xgboost")`. 
+
 # parsnip 0.1.4
 
 * `show_engines()` will provide information on the current set for a model. 

R/boost_tree.R

@@ -312,6 +312,8 @@ xgb_train <- function(
   min_child_weight = 1, gamma = 0, subsample = 1, validation = 0,
   early_stop = NULL, ...) {
 
+  others <- list(...)
+
   num_class <- length(levels(y))
 
   if (!is.numeric(validation) || validation < 0 || validation >= 1) {
@@ -327,13 +329,15 @@ xgb_train <- function(
   }
 
 
-  if (is.numeric(y)) {
-    loss <- "reg:squarederror"
-  } else {
-    if (num_class == 2) {
-      loss <- "binary:logistic"
+  if (!any(names(others) == "objective")) {
+    if (is.numeric(y)) {
+      others$objective <- "reg:squarederror"
     } else {
-      loss <- "multi:softprob"
+      if (num_class == 2) {
+        others$objective <- "binary:logistic"
+      } else {
+        others$objective <- "multi:softprob"
+      }
     }
   }
 
@@ -378,7 +382,6 @@ xgb_train <- function(
     watchlist = quote(x$watchlist),
     params = arg_list,
     nrounds = nrounds,
-    objective = loss,
     early_stopping_rounds = early_stop
   )
   if (!is.null(num_class) && num_class > 2) {
@@ -388,7 +391,7 @@ xgb_train <- function(
   call <- make_call(fun = "xgb.train", ns = "xgboost", main_args)
 
   # override or add some other args
-  others <- list(...)
+
   others <-
     others[!(names(others) %in% c("data", "weights", "nrounds", "num_class", names(arg_list)))]
   if (!(any(names(others) == "verbose"))) {

tests/testthat/test_boost_tree_xgboost.R

@@ -159,10 +159,26 @@ test_that('xgboost regression prediction', {
 
   form_pred <- predict(form_fit$fit, newdata = xgb.DMatrix(data = as.matrix(mtcars[1:8, -1])))
   expect_equal(form_pred, predict(form_fit, new_data = mtcars[1:8, -1])$.pred)
+
+  expect_equal(form_fit$fit$params$objective, "reg:squarederror")
+
 })
 
 
 
+test_that('xgboost alternate objective', {
+  skip_if_not_installed("xgboost")
+
+  spec <-
+    boost_tree() %>%
+    set_engine("xgboost", objective = "reg:pseudohubererror") %>%
+    set_mode("regression")
+
+  xgb_fit <- spec %>% fit(mpg ~ ., data = mtcars)
+  expect_equal(xgb_fit$fit$params$objective, "reg:pseudohubererror")
+})
+
+
 test_that('submodel prediction', {
 
   skip_if_not_installed("xgboost")