Better mtry check for formulas (#700)

* move dim checks for mtry and min_n to wrapper function

* better assessments of data dims

* export function and updates for multivariate cases

* pkgdown entry

Author: topepo

NAMESPACE

@@ -219,6 +219,7 @@ export(make_classes)
 export(make_engine_list)
 export(make_seealso_list)
 export(mars)
+export(max_mtry_formula)
 export(maybe_data_frame)
 export(maybe_matrix)
 export(min_cols)

R/partykit.R

@@ -90,6 +90,9 @@ cforest_train <-
     force(mtry)
     opts <- rlang::list2(...)
 
+    mtry     <- max_mtry_formula(mtry, formula, data)
+    minsplit <- min(minsplit, nrow(data))
+
     if (any(names(opts) == "control")) {
       opts$control$minsplit <- minsplit
       opts$control$maxdepth <- maxdepth
@@ -124,3 +127,24 @@ cforest_train <-
       )
     rlang::eval_tidy(forest_call)
   }
+
+# ------------------------------------------------------------------------------
+
+#' Determine largest value of mtry from formula.
+#' This function potentially caps the value of `mtry` based on a formula and
+#' data set. This is a safe approach for survival and/or multivariate models.
+#' @param mtry An initial value of `mtry` (which may be too large).
+#' @param formula A model formula.
+#' @param data The training set (data frame).
+#' @return A value for `mtry`.
+#' @examples
+#' # should be 9
+#' max_mtry_formula(200, cbind(wt, mpg) ~ ., data = mtcars)
+#' @export
+max_mtry_formula <- function(mtry, formula, data) {
+  preds <- stats::model.frame(formula, head(data))
+  trms <- attr(preds, "terms")
+  p <- ncol(attr(trms, "factors"))
+
+  max(min(mtry, p), 1L)
+}

R/rand_forest.R

@@ -163,21 +163,15 @@ translate.rand_forest <- function(x, engine = x$engine, ...) {
   ## -----------------------------------------------------------------------------
   # Protect some arguments based on data dimensions
 
-  if (any(names(arg_vals) == "mtry") & engine != "cforest") {
+  if (any(names(arg_vals) == "mtry") & engine != "partykit") {
     arg_vals$mtry <- rlang::call2("min_cols", arg_vals$mtry, expr(x))
   }
-  if (any(names(arg_vals) == "mtry") & engine == "cforest") {
-    arg_vals$mtry <- rlang::call2("min_cols", arg_vals$mtry, expr(data))
-  }
 
   if (any(names(arg_vals) == "min.node.size")) {
     arg_vals$min.node.size <-
       rlang::call2("min_rows", arg_vals$min.node.size, expr(x))
   }
-  if (any(names(arg_vals) == "minsplit" & engine == "cforest")) {
-    arg_vals$minsplit <-
-      rlang::call2("min_rows", arg_vals$minsplit, expr(data))
-  }
+
   if (any(names(arg_vals) == "nodesize")) {
     arg_vals$nodesize <-
       rlang::call2("min_rows", arg_vals$nodesize, expr(x))

_pkgdown.yml

@@ -95,6 +95,7 @@ reference:
       - set_new_model
       - maybe_matrix
       - min_cols
+      - max_mtry_formula
       - required_pkgs
       - required_pkgs.model_spec
       - req_pkgs

man/max_mtry_formula.Rd

@@ -83,4 +83,23 @@ test_that('control class', {
   )
 })
 
+# ------------------------------------------------------------------------------
+
+test_that('correct mtry', {
+  skip_if_not_installed("modeldata")
+  data(ames, package = "modeldata")
+  f_1 <- Sale_Price ~ Longitude + Latitude + Year_Built
+  f_2 <- Sale_Price ~ .
+  f_3 <- cbind(wt, mpg) ~ .
+
+  expect_equal(max_mtry_formula(2, f_1, ames), 2)
+  expect_equal(max_mtry_formula(5, f_1, ames), 3)
+  expect_equal(max_mtry_formula(0, f_1, ames), 1)
+
+  expect_equal(max_mtry_formula(2000, f_2, ames), ncol(ames) - 1)
+  expect_equal(max_mtry_formula(2, f_2, ames), 2)
+
+  expect_equal(max_mtry_formula(200, f_3, data = mtcars), ncol(mtcars) - 2)
+
+})