R/cv_param_grid.R

In nanxstats/stackgbm: Stacked Gradient Boosting Machines

#' Generate a parameter grid for cross-validation
#'
#' This function generates a parameter grid to be used in the
#' cross-validation of gradient boosting decision tree (GBDT) models.
#'
#' @param n_iterations A numeric vector of the number of iterations (trees)
#'   for the GBDT model. This is equivalent to `nrounds` in XGBoost,
#'   `num_iterations` in LightGBM, and `iterations` in CatBoost.
#' @param max_depth A numeric vector of the maximum tree depths.
#'   This parameter is equivalent to `max_depth` in XGBoost and LightGBM,
#'   and `depth` in CatBoost.
#' @param learning_rate A numeric vector of learning rates for the GBDT model.
#'   This parameter is equivalent to `eta` in XGBoost,
#'   `learning_rate` in LightGBM, and ignored in CatBoost.
#'
#' @return
#' A list where the names are the parameter names and the values
#' are vectors of possible values for those parameters.
#'
#' @export
#'
#' @examples
#' params <- cv_param_grid(
#'   n_iterations = c(10, 100),
#'   max_depth = c(3, 5),
#'   learning_rate = c(0.01, 0.1)
#' )
cv_param_grid <- function(
    n_iterations = c(100, 200, 500, 1000),
    max_depth = c(3, 5, 7, 9),
    learning_rate = c(0.01, 0.05, 0.1, 0.2)) {
  list(
    "n_iterations" = n_iterations,
    "max_depth" = max_depth,
    "learning_rate" = learning_rate
  )
}

#' Map stackgbm parameter names to xgboost parameter names
#'
#' @param params Parameter grid generated by [cv_param_grid()].
#'
#' @return Mapped list of parameter grid.
#'
#' @noRd
map_params_xgboost <- function(params) {
  list(
    nrounds = params$n_iterations,
    max_depth = params$max_depth,
    eta = params$learning_rate
  )
}

#' Map stackgbm parameter names to lightgbm parameter names
#'
#' @param params Parameter grid generated by [cv_param_grid()].
#'
#' @return Mapped list of parameter grid.
#'
#' @noRd
map_params_lightgbm <- function(params) {
  list(
    num_iterations = params$n_iterations,
    max_depth = params$max_depth,
    learning_rate = params$learning_rate
  )
}

#' Map stackgbm parameter names to lightgbm parameter names
#'
#' @param params Parameter grid generated by [cv_param_grid()].
#'
#' @return Mapped list of parameter grid.
#'
#' @noRd
map_params_catboost <- function(params) {
  list(
    iterations = params$n_iterations,
    depth = params$max_depth
  )
}