Add option to keep resamples by fold

- Add aggregate_resamples parameter to extractBestPreds
- Pass parameter through caretList chain
- Add unit tests for new functionality
- Maintain backward compatibility with default TRUE

R/caretList.R

@@ -18,6 +18,7 @@ table #' Create a list of several train models from the caret package
 #' @param continue_on_fail logical, should a valid caretList be returned that
 #' excludes models that fail, default is FALSE
 #' @param trim logical should the train models be trimmed to save memory and speed up stacking
+#' @param aggregate_resamples logical, whether to aggregate resamples by taking mean/first value. Default TRUE.
 #' @return A list of \code{\link[caret]{train}} objects. If the model fails to build,
 #' it is dropped from the list.
 #' @export
@@ -37,7 +38,8 @@ caretList <- function(
     tuneList = NULL,
     metric = NULL,
     continue_on_fail = FALSE,
-    trim = TRUE) {
+    trim = TRUE,
+    aggregate_resamples = TRUE) {
   # Checks
   if (is.null(tuneList) && is.null(methodList)) {
     stop("Please either define a methodList or tuneList", call. = FALSE)
@@ -79,7 +81,8 @@ caretList <- function(
   global_args[["metric"]] <- metric
 
   # Loop through the tuneLists and fit caret models with those specs
-  modelList <- lapply(tuneList, caretTrain, global_args = global_args, continue_on_fail = continue_on_fail, trim = trim)
+  modelList <- lapply(tuneList, caretTrain, global_args = global_args, continue_on_fail = continue_on_fail, 
+                     trim = trim, aggregate_resamples = aggregate_resamples)
   names(modelList) <- names(tuneList)
   nulls <- vapply(modelList, is.null, logical(1L))
   modelList <- modelList[!nulls]

R/caretPredict.R

@@ -80,9 +80,10 @@ caretPredict <- function(object, newdata = NULL, excluded_class_id = 1L, ...) {
 #'  If `TRUE`, the function will return `NULL` if the `train` function fails.
 #' @param trim A logical indicating whether to trim the output model.
 #' If `TRUE`, the function will remove some elements that are not needed from the output model.
+#' @param aggregate_resamples A logical indicating whether to aggregate resamples by taking mean/first value.
 #' @return The output of the `train` function.
 #' @keywords internal
-caretTrain <- function(local_args, global_args, continue_on_fail = FALSE, trim = TRUE) {
+caretTrain <- function(local_args, global_args, continue_on_fail = FALSE, trim = TRUE, aggregate_resamples = TRUE) {
   # Combine args
   # I think my handling here is correct (update globals with locals, which allows locals be partial)
   # but it would be nice to have some tests
@@ -100,7 +101,7 @@ caretTrain <- function(local_args, global_args, continue_on_fail = FALSE, trim =
 
   # Only save stacked predictions for the best model
   if ("pred" %in% names(model)) {
-    model[["pred"]] <- extractBestPreds(model)
+    model[["pred"]] <- extractBestPreds(model, aggregate_resamples = aggregate_resamples)
   }
 
   if (trim) {
@@ -147,9 +148,10 @@ aggregate_mean_or_first <- function(x) {
 #' @title Extract the best predictions from a train object
 #' @description Extract the best predictions from a train object.
 #' @param x a train object
+#' @param aggregate_resamples logical, whether to aggregate resamples by taking mean/first value. Default TRUE.
 #' @return a data.table::data.table with predictions
 #' @keywords internal
-extractBestPreds <- function(x) {
+extractBestPreds <- function(x, aggregate_resamples = TRUE) {
   stopifnot(methods::is(x, "train"))
   if (is.null(x$pred)) {
     stop("No predictions saved during training. Please set savePredictions = 'final' in trainControl", call. = FALSE)
@@ -167,14 +169,19 @@ extractBestPreds <- function(x) {
   # Drop rows for other tunes
   pred <- pred[best_tune, ]
 
-  # If we have multiple resamples per row
+  # If we have multiple resamples per row and aggregate_resamples is TRUE
   # e.g. for repeated CV, we need to average the predictions
-  keys <- "rowIndex"
-  data.table::setkeyv(pred, keys)
-  pred <- pred[, lapply(.SD, aggregate_mean_or_first), by = keys]
+  if (aggregate_resamples) {
+    keys <- "rowIndex"
+    data.table::setkeyv(pred, keys)
+    pred <- pred[, lapply(.SD, aggregate_mean_or_first), by = keys]
+  } else {
+    # When not aggregating, keep all columns but ensure consistent order
+    data.table::setcolorder(pred, c("rowIndex", setdiff(names(pred), "rowIndex")))
+  }
 
   # Order results consistently
-  data.table::setorderv(pred, keys)
+  data.table::setorderv(pred, "rowIndex")
 
   # Return
   pred

tests/testthat/test-caretList.R

@@ -146,6 +146,48 @@ testthat::test_that("predict.caretList works for classification and regression",
 testthat::context("caretList")
 ################################################################
 
+testthat::test_that("caretList respects aggregate_resamples parameter", {
+  # Create a trainControl with repeated CV to ensure multiple resamples per row
+  ctrl <- caret::trainControl(
+    method = "repeatedcv",
+    number = 3,
+    repeats = 2,
+    savePredictions = "final"
+  )
+
+  # Test with aggregate_resamples = FALSE
+  test_no_agg <- caretList(
+    x = train[, -23L],
+    y = train[, "Class"],
+    methodList = c("knn"),
+    trControl = ctrl,
+    aggregate_resamples = FALSE
+  )
+
+  # Get predictions from the first model
+  preds_no_agg <- test_no_agg[[1]]$pred
+
+  # Test with aggregate_resamples = TRUE (default)
+  test_agg <- caretList(
+    x = train[, -23L],
+    y = train[, "Class"],
+    methodList = c("knn"),
+    trControl = ctrl
+  )
+
+  # Get predictions from the first model
+  preds_agg <- test_agg[[1]]$pred
+
+  # The non-aggregated predictions should have more rows than aggregated ones
+  # since they keep all resamples
+  testthat::expect_gt(nrow(preds_no_agg), nrow(preds_agg))
+
+  # Both should have rowIndex as first column and same set of columns
+  testthat::expect_identical(names(preds_no_agg)[1], "rowIndex")
+  testthat::expect_identical(names(preds_agg)[1], "rowIndex")
+  testthat::expect_setequal(names(preds_no_agg), names(preds_agg))
+})
+
 testthat::test_that("caretList works for various scenarios", {
   # Basic classification
   test1 <- caretList(