retain .iter in collector() (#666)

* retain `.iter` in `collector()` * note in NEWS
tidymodels · Apr 11, 2023 · e0fbefd · e0fbefd
1 parent fd7c3d9
commit e0fbefd
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,8 @@
 # tune (development version)
 
+* Fixed a bug introduced in tune 1.1.0 in `collect_()` functions where the 
+  `.iter` column was dropped.
+
 # tune 1.1.0
 
 tune 1.1.0 introduces a number of new features and bug fixes, accompanied by various optimizations that substantially decrease the total evaluation time to tune hyperparameters in the tidymodels.

diff --git a/R/collect.R b/R/collect.R
@@ -361,27 +361,37 @@ collect_metrics.tune_results <- function(x, summarize = TRUE, ...) {
 }
 
 collector <- function(x, coll_col = ".predictions") {
-  if (any(colnames(x) == ".iter")) {
+  is_iterative <- any(colnames(x) == ".iter")
+  if (is_iterative) {
     keep_cols <- c(coll_col, ".iter")
   } else {
     keep_cols <- coll_col
   }
 
   id_cols <- colnames(x)[grepl("id", colnames(x))]
-  keep_cols <- c(coll_col, id_cols)
+  keep_cols <- c(id_cols, keep_cols)
   x <- x[keep_cols]
   coll_col <- x[[coll_col]]
+  sizes <- vctrs::list_sizes(coll_col)
 
   res <-
     vctrs::vec_cbind(
-      vctrs::vec_rep_each(x[, id_cols], times = vctrs::list_sizes(coll_col)),
+      vctrs::vec_rep_each(x[, id_cols], times = sizes),
       vctrs::list_unchop(coll_col)
     )
 
+  if (is_iterative) {
+    res <-
+      vctrs::vec_cbind(
+        res,
+        vctrs::vec_rep_each(x[, ".iter"], times = sizes)
+      )
+  }
+
   arrange_cols <- c(".iter", ".config")
   arrange_cols <- arrange_cols[rlang::has_name(res, arrange_cols)]
 
-  res <- vctrs::vec_slice(res, vctrs::vec_order(res[arrange_cols]))
+  vctrs::vec_slice(res, vctrs::vec_order(res[arrange_cols]))
 }
 
 #' @export

diff --git a/tests/testthat/test-collect.R b/tests/testthat/test-collect.R
@@ -112,6 +112,11 @@ test_that("classification class predictions, averaged", {
   res <- collect_predictions(svm_tune_class, summarize = TRUE)
   expect_equal(nrow(res), nrow(two_class_dat) * nrow(svm_grd))
   expect_false(dplyr::is_grouped_df(res))
+  expect_named(
+    collect_predictions(svm_tune, summarize = TRUE),
+    c(".row", "cost value", "Class", ".config", ".iter", ".pred_Class1",
+      ".pred_Class2", ".pred_class")
+  )
 
   # pull out an example to test
   all_res_subset <-