handle "class"-like outcome names and levels

R/add_candidates.R

@@ -279,10 +279,9 @@ add_candidates.default <- function(data_stack, candidates, name, ...) {
     ) %>%
     dplyr::select(-.row) 
 
-  pred_class_idx <- grepl(pattern = ".pred_class", x = colnames(candidate_cols))
-
-  candidate_cols <- candidate_cols[,!pred_class_idx] %>% 
-    setNames(., make.names(names(.)))
+  if (attr(stack, "mode") == "classification") {
+    candidate_cols <- remove_class_preds(candidate_cols)
+  }
 
   if (nrow(stack) == 0) {
     stack <- 
@@ -502,3 +501,23 @@ collate_predictions <- function(x) {
   }
   res
 }
+
+# given a set of candidate columns, removes those with hard class predictions
+remove_class_preds <- function(x) {
+  lvls <- levels(factor(x[[1]]))
+
+  # gather indices for the columns with class probability predictions
+  prob_preds_idx <- purrr::map(paste0(".pred_", lvls), grepl, x = colnames(x)) %>%
+    purrr::pmap(any) %>%
+    unlist()
+
+  # select the columns that look like they have probability predictions,
+  # get rid of entries that would have been okayed because of an outcome
+  # level called "class", and re-attach the outcome column itself
+  x[,prob_preds_idx] %>%
+    dplyr::select(where(is.numeric)) %>% 
+    dplyr::bind_cols(x[,1], .) %>%
+    setNames(., make.names(names(.)))
+}
+
+

tests/testthat/test_add_candidates.R

@@ -261,3 +261,29 @@ test_that("model definition naming works as expected", {
     "cannot prefix a valid column name"
   )
 })
+
+test_that("stacks can handle columns and levels named 'class'", {
+  # waiting on https://github.com/tidymodels/tune/issues/487
+  # to be able to test with entry "class"
+  x <- tibble::tibble(
+    class = sample(c("class_1", "class_2"), 100, replace = TRUE),
+    a = rnorm(100),
+    b = rnorm(100)
+  )
+
+  res <- tune::tune_grid(
+    parsnip::logistic_reg(engine = 'glmnet', penalty = tune::tune(), mixture = 1),
+    preprocessor = recipes::recipe(class ~ ., x),
+    resamples = rsample::vfold_cv(x, 2),
+    grid = 2,
+    control = control_stack_grid()
+  )
+
+  expect_s3_class(
+    suppressWarnings(
+      stacks() %>% 
+        add_candidates(res)
+    ),
+    "data_stack"
+  )
+})