protype changes for #184

Author: topepo

R/rand_forest.R

@@ -161,6 +161,8 @@ translate.rand_forest <- function(x, engine = x$engine, ...) {
 
   x <- translate.default(x, engine, ...)
 
+  ## -----------------------------------------------------------------------------
+
   # slightly cleaner code using
   arg_vals <- x$method$fit$args
 
@@ -185,14 +187,40 @@ translate.rand_forest <- function(x, engine = x$engine, ...) {
 
   # add checks to error trap or change things for this method
   if (engine == "ranger") {
-    if (any(names(arg_vals) == "importance"))
-      if (isTRUE(is.logical(quo_get_expr(arg_vals$importance))))
+
+    if (any(names(arg_vals) == "importance")) {
+      if (isTRUE(is.logical(quo_get_expr(arg_vals$importance)))) {
         rlang::abort("`importance` should be a character value. See ?ranger::ranger.")
+      }
+    }
     # unless otherwise specified, classification models are probability forests
-    if (x$mode == "classification" && !any(names(arg_vals) == "probability"))
+    if (x$mode == "classification" && !any(names(arg_vals) == "probability")) {
       arg_vals$probability <- TRUE
+    }
+  }
+
+  ## -----------------------------------------------------------------------------
+  # Protect some arguments based on data dimensions
 
+  if (any(names(arg_vals) == "mtry")) {
+    arg_vals$mtry <- rlang::call2("min", arg_vals$mtry, expr(ncol(x)))
   }
+
+  if (any(names(arg_vals) == "min.node.size")) {
+    arg_vals$min.node.size <-
+      rlang::call2("min", arg_vals$min.node.size, expr(nrow(x)))
+  }
+  if (any(names(arg_vals) == "nodesize")) {
+    arg_vals$nodesize <-
+      rlang::call2("min", arg_vals$nodesize, expr(nrow(x)))
+  }
+  if (any(names(arg_vals) == "min_instances_per_node")) {
+    arg_vals$min_instances_per_node <-
+      rlang::call2("min", arg_vals$min_instances_per_node, expr(nrow(x)))
+  }
+
+  ## -----------------------------------------------------------------------------
+
   x$method$fit$args <- arg_vals
 
   x

tests/testthat/test_rand_forest_randomForest.R

@@ -213,3 +213,29 @@ test_that('randomForest regression prediction', {
   expect_equal(xy_pred, predict(xy_fit, new_data = tail(mtcars))$.pred)
 
 })
+
+## -----------------------------------------------------------------------------
+
+test_that('argument checks for data dimensions', {
+
+  skip_if_not_installed("randomForest")
+
+  data(penguins, package = "modeldata")
+  penguins <- na.omit(penguins)
+
+  spec <-
+    rand_forest(mtry = 1000, min_n = 1000, trees = 5) %>%
+    set_engine("randomForest") %>%
+    set_mode("regression")
+
+  f_fit  <- spec %>% fit(body_mass_g ~ ., data = penguins)
+  xy_fit <- spec %>% fit_xy(x = penguins[, -6], y = penguins$body_mass_g)
+
+  expect_equal(f_fit$fit$mtry, 6)
+  expect_equal(f_fit$fit$call$nodesize, rlang::expr(min(~1000, nrow(x))))
+  expect_equal(xy_fit$fit$mtry, 6)
+  expect_equal(xy_fit$fit$call$nodesize, rlang::expr(min(~1000, nrow(x))))
+
+})
+
+

tests/testthat/test_rand_forest_ranger.R

@@ -474,3 +474,27 @@ test_that('ranger and sparse matrices', {
 
 })
 
+
+## -----------------------------------------------------------------------------
+
+test_that('argument checks for data dimensions', {
+
+  skip_if_not_installed("ranger")
+
+  data(penguins, package = "modeldata")
+  penguins <- na.omit(penguins)
+
+  spec <-
+    rand_forest(mtry = 1000, min_n = 1000, trees = 5) %>%
+    set_engine("ranger") %>%
+    set_mode("regression")
+
+  f_fit  <- spec %>% fit(body_mass_g ~ ., data = penguins)
+  xy_fit <- spec %>% fit_xy(x = penguins[, -6], y = penguins$body_mass_g)
+
+  expect_equal(f_fit$fit$mtry, 6)
+  expect_equal(f_fit$fit$min.node.size, nrow(penguins))
+  expect_equal(xy_fit$fit$mtry, 6)
+  expect_equal(xy_fit$fit$min.node.size, nrow(penguins))
+
+})