Add Fold Weights for Variable Resample Weighting

NAMESPACE

@@ -82,7 +82,9 @@ S3method(pretty,tune_results)
 S3method(print,control_bayes)
 S3method(print,control_grid)
 S3method(print,control_last_fit)
+S3method(print,manual_rset)
 S3method(print,prob_improve)
+S3method(print,rset)
 S3method(print,tune_results)
 S3method(select_best,default)
 S3method(select_best,tune_results)
@@ -149,8 +151,10 @@ export(.get_tune_workflow)
 export(.load_namespace)
 export(.stash_last_result)
 export(.use_case_weights_with_yardstick)
+export(add_fold_weights)
 export(augment)
 export(autoplot)
+export(calculate_fold_weights)
 export(check_eval_time_arg)
 export(check_initial)
 export(check_metric_in_tune_results)
@@ -203,6 +207,7 @@ export(fit_resamples)
 export(forge_from_workflow)
 export(future_installed)
 export(get_future_workers)
+export(get_fold_weights)
 export(get_metric_time)
 export(get_mirai_workers)
 export(get_parallel_seeds)

R/checks.R

@@ -22,6 +22,26 @@ check_rset <- function(x) {
   if (inherits(x, "permutations")) {
     cli::cli_abort("Permutation samples are not suitable for tuning.")
   }
+
+  # Check fold weights if present
+  check_fold_weights(x)
+
+  invisible(NULL)
+}
+
+#' Check fold weights in rset objects
+#'
+#' @param x An rset object.
+#' @return `NULL` invisibly, or error if weights are invalid.
+#' @keywords internal
+check_fold_weights <- function(x) {
+  weights <- attr(x, ".fold_weights")
+  if (is.null(weights)) {
+    return(invisible(NULL))
+  }
+
+  .validate_fold_weights(weights, nrow(x))
+
   invisible(NULL)
 }
 

R/collect.R

@@ -565,6 +565,8 @@ estimate_tune_results <- function(x, ..., col_name = ".metrics") {
     )
   }
 
+  fold_weights <- .get_fold_weights(x)
+
   # The mapping of tuning parameters and .config.
   config_key <- .config_key_from_metrics(x)
 
@@ -604,19 +606,51 @@ estimate_tune_results <- function(x, ..., col_name = ".metrics") {
 
   x <- tibble::as_tibble(x)
   x <- vctrs::vec_slice(x, x$id != "Apparent")
-  x <- x |>
-    dplyr::group_by(
-      !!!rlang::syms(param_names),
-      .metric,
-      .estimator,
-      !!!rlang::syms(group_cols)
-    ) |>
-    dplyr::summarize(
-      mean = mean(.estimate, na.rm = TRUE),
-      n = sum(!is.na(.estimate)),
-      std_err = sd(.estimate, na.rm = TRUE) / sqrt(n),
-      .groups = "drop"
-    )
+
+  # Join weights to the data if available
+  if (!is.null(fold_weights)) {
+    weight_data <- .create_weight_mapping(fold_weights, id_names, x)
+    if (!is.null(weight_data)) {
+      x <- dplyr::left_join(x, weight_data, by = id_names)
+    } else {
+      # If weight mapping failed, fall back to unweighted
+      fold_weights <- NULL
+    }
+  }
+
+  if (!is.null(fold_weights)) {
+    # Use weighted aggregation
+    x <- x |>
+      dplyr::group_by(
+        !!!rlang::syms(param_names),
+        .metric,
+        .estimator,
+        !!!rlang::syms(group_cols)
+      ) |>
+      dplyr::summarize(
+        mean = .weighted_mean(.estimate, .fold_weight),
+        n = sum(!is.na(.estimate)),
+        effective_n = .effective_sample_size(.fold_weight[!is.na(.estimate)]),
+        std_err = .weighted_sd(.estimate, .fold_weight) /
+          sqrt(pmax(effective_n, 1)),
+        .groups = "drop"
+      ) |>
+      dplyr::select(-effective_n)
+  } else {
+    x <- x |>
+      dplyr::group_by(
+        !!!rlang::syms(param_names),
+        .metric,
+        .estimator,
+        !!!rlang::syms(group_cols)
+      ) |>
+      dplyr::summarize(
+        mean = mean(.estimate, na.rm = TRUE),
+        n = sum(!is.na(.estimate)),
+        std_err = sd(.estimate, na.rm = TRUE) / sqrt(n),
+        .groups = "drop"
+      )
+  }
 
   # only join when parameters are being tuned (#600)
   if (length(param_names) == 0) {

R/tune_grid.R

@@ -425,6 +425,10 @@ pull_rset_attributes <- function(x) {
   att$class <- setdiff(class(x), class(tibble::new_tibble(list())))
   att$class <- att$class[att$class != "rset"]
 
+  if (!is.null(attr(x, ".fold_weights"))) {
+    att[[".fold_weights"]] <- attr(x, ".fold_weights")
+  }
+
   lab <- try(pretty(x), silent = TRUE)
   if (inherits(lab, "try-error")) {
     lab <- NA_character_