enable xgboost to use sparse X for tidymodels/tidymodels#42

Author: topepo

DESCRIPTION

@@ -51,4 +51,5 @@ Suggests:
     MASS,
     nlme,
     modeldata,
-    liquidSVM
+    liquidSVM,
+    Matrix

R/boost_tree.R

@@ -290,11 +290,8 @@ xgb_train <- function(
   min_child_weight = 1, gamma = 0, subsample = 1, validation = 0,
   early_stop = NULL, ...) {
 
-  if (length(levels(y)) > 2) {
-    num_class <- length(levels(y))
-  }  else {
-    num_class <- NULL
-  }
+  num_class <- length(levels(y))
+
   if (!is.numeric(validation) || validation < 0 || validation >= 1) {
     rlang::abort("`validation` should be on [0, 1).")
   }
@@ -311,36 +308,17 @@ xgb_train <- function(
   if (is.numeric(y)) {
     loss <- "reg:squarederror"
   } else {
-    lvl <- levels(y)
-    y <- as.numeric(y) - 1
-    if (length(lvl) == 2) {
+    if (num_class == 2) {
       loss <- "binary:logistic"
     } else {
       loss <- "multi:softprob"
     }
   }
 
-  if (is.data.frame(x)) {
-    x <- as.matrix(x) # maybe use model.matrix here?
-  }
-
   n <- nrow(x)
   p <- ncol(x)
 
-  if (!inherits(x, "xgb.DMatrix")) {
-    if (validation > 0) {
-      trn_index <- sample(1:n, size = floor(n * validation) + 1)
-      wlist <-
-        list(validation = xgboost::xgb.DMatrix(x[-trn_index, ], label = y[-trn_index], missing = NA))
-      x <- xgboost::xgb.DMatrix(x[trn_index, ], label = y[trn_index], missing = NA)
-
-    } else {
-      x <- xgboost::xgb.DMatrix(x, label = y, missing = NA)
-      wlist <- list(training = x)
-    }
-  } else {
-    xgboost::setinfo(x, "label", y)
-  }
+  x <- as_xgb_data(x, y, validation)
 
   # translate `subsample` and `colsample_bytree` to be on (0, 1] if not
   if (subsample > 1) {
@@ -366,17 +344,15 @@ xgb_train <- function(
     subsample = subsample
   )
 
-  # eval if contains expressions?
-
   main_args <- list(
-    data = quote(x),
-    watchlist = quote(wlist),
+    data = quote(x$data),
+    watchlist = quote(x$watchlist),
     params = arg_list,
     nrounds = nrounds,
     objective = loss,
     early_stopping_rounds = early_stop
   )
-  if (!is.null(num_class)) {
+  if (!is.null(num_class) && num_class > 2) {
     main_args$num_class <- num_class
   }
 
@@ -399,7 +375,7 @@ xgb_train <- function(
 #' @importFrom stats binomial
 xgb_pred <- function(object, newdata, ...) {
   if (!inherits(newdata, "xgb.DMatrix")) {
-    newdata <- as.matrix(newdata)
+    newdata <- as_matrix(newdata)
     newdata <- xgboost::xgb.DMatrix(data = newdata, missing = NA)
   }
 
@@ -415,6 +391,37 @@ xgb_pred <- function(object, newdata, ...) {
   x
 }
 
+
+as_xgb_data <- function(x, y, validation = 0, ...) {
+  lvls <- levels(y)
+  n <- nrow(x)
+
+  if (is.data.frame(x)) {
+    x <- as.matrix(x)
+  }
+
+  if (is.factor(y)) {
+    y <- as.numeric(y) - 1
+  }
+
+  if (!inherits(x, "xgb.DMatrix")) {
+    if (validation > 0) {
+      trn_index <- sample(1:n, size = floor(n * (1 - validation)) + 1)
+      wlist <-
+        list(validation = xgboost::xgb.DMatrix(x[-trn_index, ], label = y[-trn_index], missing = NA))
+      dat <- xgboost::xgb.DMatrix(x[trn_index, ], label = y[trn_index], missing = NA)
+
+    } else {
+      dat <- xgboost::xgb.DMatrix(x, label = y, missing = NA)
+      wlist <- list(training = dat)
+    }
+  } else {
+    dat <- xgboost::setinfo(x, "label", y)
+    wlist <- list(training = dat)
+  }
+
+  list(data = dat, watchlist = wlist)
+}
 #' @importFrom purrr map_df
 #' @export
 #' @rdname multi_predict

tests/testthat/test_boost_tree_xgboost.R

@@ -286,3 +286,70 @@ test_that('early stopping', {
    regex = "`early_stop` should be on"
  )
 })
+
+
+## -----------------------------------------------------------------------------
+
+test_that('xgboost data conversion', {
+  skip_if_not_installed("xgboost")
+
+  mtcar_x <- mtcars[, -1]
+  mtcar_mat <- as.matrix(mtcar_x)
+  mtcar_smat <- Matrix::Matrix(mtcar_mat, sparse = TRUE)
+
+  expect_error(from_df <- parsnip:::as_xgb_data(mtcar_x, mtcars$mpg), regexp = NA)
+  expect_true(inherits(from_df$data, "xgb.DMatrix"))
+  expect_true(inherits(from_df$watchlist$training, "xgb.DMatrix"))
+
+  expect_error(from_mat <- parsnip:::as_xgb_data(mtcar_mat, mtcars$mpg), regexp = NA)
+  expect_true(inherits(from_mat$data, "xgb.DMatrix"))
+  expect_true(inherits(from_mat$watchlist$training, "xgb.DMatrix"))
+
+  expect_error(from_sparse <- parsnip:::as_xgb_data(mtcar_smat, mtcars$mpg), regexp = NA)
+  expect_true(inherits(from_mat$data, "xgb.DMatrix"))
+  expect_true(inherits(from_mat$watchlist$training, "xgb.DMatrix"))
+
+  expect_error(from_df <- parsnip:::as_xgb_data(mtcar_x, mtcars$mpg, validation = .1), regexp = NA)
+  expect_true(inherits(from_df$data, "xgb.DMatrix"))
+  expect_true(inherits(from_df$watchlist$validation, "xgb.DMatrix"))
+  expect_true(nrow(from_df$data) > nrow(from_df$watchlist$validation))
+
+  expect_error(from_mat <- parsnip:::as_xgb_data(mtcar_mat, mtcars$mpg, validation = .1), regexp = NA)
+  expect_true(inherits(from_mat$data, "xgb.DMatrix"))
+  expect_true(inherits(from_mat$watchlist$validation, "xgb.DMatrix"))
+  expect_true(nrow(from_mat$data) > nrow(from_mat$watchlist$validation))
+
+  expect_error(from_sparse <- parsnip:::as_xgb_data(mtcar_smat, mtcars$mpg, validation = .1), regexp = NA)
+  expect_true(inherits(from_mat$data, "xgb.DMatrix"))
+  expect_true(inherits(from_mat$watchlist$validation, "xgb.DMatrix"))
+  expect_true(nrow(from_sparse$data) > nrow(from_sparse$watchlist$validation))
+
+})
+
+
+test_that('xgboost data and sparse matrices', {
+  skip_if_not_installed("xgboost")
+
+  mtcar_x <- mtcars[, -1]
+  mtcar_mat <- as.matrix(mtcar_x)
+  mtcar_smat <- Matrix::Matrix(mtcar_mat, sparse = TRUE)
+
+  xgb_spec <-
+    boost_tree(trees = 10) %>%
+    set_engine("xgboost") %>%
+    set_mode("regression")
+
+  set.seed(1)
+  from_df <- xgb_spec %>% fit_xy(mtcar_x, mtcars$mpg)
+  set.seed(1)
+  from_mat <- xgb_spec %>% fit_xy(mtcar_mat, mtcars$mpg)
+  set.seed(1)
+  from_sparse <- xgb_spec %>% fit_xy(mtcar_smat, mtcars$mpg)
+
+  expect_equal(from_df$fit, from_mat$fit)
+  expect_equal(from_df$fit, from_sparse$fit)
+
+})
+
+
+

tests/testthat/test_convert_data.R

@@ -626,4 +626,8 @@ test_that("convert to matrix", {
   expect_true(inherits(parsnip::as_matrix(mtcars), "matrix"))
   expect_true(inherits(parsnip::as_matrix(tibble::as_tibble(mtcars)), "matrix"))
   expect_true(inherits(parsnip::as_matrix(as.matrix(mtcars)), "matrix"))
+  expect_true(
+    inherits(parsnip::as_matrix(Matrix::Matrix(as.matrix(mtcars), sparse = TRUE)),
+             "dgCMatrix")
+  )
 })