R/xgboost.model.R

In rladeira/stacking:

xgboost.model <- function(tune.with.cv = FALSE,
                          nfold = 5,
                          cv.max.nrounds = 10000,
                          early.stop.round = 10,
                          stratified = TRUE,
                          maximize = FALSE,
                          preProcess = NULL,
                          description = NULL,
                          ...)
{
  library(xgboost)

  function()
  {
    model.name <- "xgboost"
    model_ <- NULL

    train_ <- function(X_train, y)
    {
      y <- as.numeric(y) - 1

      if (tune.with.cv)
      {
        cvResult <- xgb.cv(data = as.matrix(X_train),
                           label = y,
                           nfold = nfold,
                           nrounds = cv.max.nrounds,
                           early.stop.round = early.stop.round,
                           stratified = stratified,
                           maximize = maximize,
                           ...)
        cat("\n")

        model_ <<- xgboost(data = as.matrix(X_train),
                           label = y,
                           nrounds = ifelse(maximize,
                                            yes = which.max(cvResult[[3]]),
                                            no  = which.min(cvResult[[3]])),
                           ...)
      } else {
        model_ <<- xgboost(data = as.matrix(X_train),
                           label = y,  ...)
      }

      invisible()
    }

    predict_ <- function(X_test)
    {
      predictions <- predict(model_, X_test)
      predictions <- matrix(predictions, nrow = nrow(X_test), byrow = TRUE)
      colnames(predictions) <- paste(model.name, 1:ncol(predictions), sep = "_")
      predictions
    }

    list(
      train_ = train_,
      predict_ = predict_,
      name = model.name,
      description = description
    )
  }
}