R/split_cv.R

In mintyr: Grouped and Nested Data Pipelines Built on 'data.table'

# WARNING - Generated by {fusen} from dev/flat_cross_validation.Rmd: do not edit by hand # nolint: line_length_linter.

#' Apply V-Fold Cross-Validation to a List of Datasets
#'
#' @description
#' `split_cv` creates (optionally repeated and stratified) v-fold
#' cross-validation splits for every dataset in a list. It is implemented
#' with base R and `data.table` only and returns, per dataset, a
#' `data.table` of fold identifiers, row indices and (optionally) the
#' training / validation subsets.
#'
#' @param data     A `list` whose every element is a `data.frame` or
#'   `data.table`. Must be non-empty. A single `data.frame` is not accepted
#'   (wrap it in `list()`).
#' @param v        Number of folds. A single integer >= 2. Default `10`.
#' @param repeats  Number of repeats. A single integer >= 1. Default `1`.
#' @param strata   `NULL` (default) or a single column name present in
#'   **every** dataset. Rows are shuffled within each stratum so that every
#'   fold receives a proportional share of each stratum.
#' @param breaks   Number of quantile bins used when `strata` is numeric with
#'   more than `breaks` distinct values. Default `4`.
#' @param pool     Strata holding less than this proportion of the rows are
#'   merged with the next-smallest stratum. A number in `[0, 0.5)`.
#'   Default `0.1`.
#' @param seed     `NULL` (default) or a single number for reproducible folds.
#'   The caller's random-number stream is restored on exit.
#' @param materialize Logical. If `TRUE` (default), `train` / `validate`
#'   list-columns holding the subsets are added. `FALSE` keeps only the row
#'   indices, which needs far less memory for large data.
#' @param out_type Class of the `train` / `validate` subsets: `"dt"`
#'   (`data.table`, default) or `"df"` (plain `data.frame`, e.g. for
#'   'ASReml-R' or other modelling functions that expect a `data.frame`).
#'   Factor levels are kept in full in every subset. Only used when
#'   `materialize = TRUE`; the returned fold table itself is always a
#'   `data.table`.
#'
#' @details
#' Fold sizes differ by at most one row. With stratification, rows are
#' shuffled within strata, the strata are laid out one after another and fold
#' labels are dealt out cyclically, so each stratum is spread evenly across
#' folds. The stratification rules (`breaks`, `pool`) follow the same idea
#' as `rsample::vfold_cv()`, but the random fold assignments are not
#' identical to rsample's.
#'
#' Training and validation subsets can always be rebuilt from the indices:
#' `data[[1]][res[[1]]$train_idx[[k]], ]`.
#'
#' For mixed-model cross-validation (e.g. genomic prediction with
#' 'ASReml-R') a common alternative to subsetting is masking: keep all
#' records, set the phenotypes of `validate_idx` to `NA`, fit the model on
#' the full data and correlate the predicted breeding values of the masked
#' individuals with their observations. `materialize = FALSE` is enough for
#' this, since only the indices are needed.
#'
#' @return A `list` of `data.table` objects (one per input dataset, names of
#'   `data` preserved), each with one row per fold and the columns:
#' \itemize{
#'   \item `id`           — fold label (`Fold1`, ...), or repeat label
#'         (`Repeat1`, ...) when `repeats > 1`.
#'   \item `id2`          — fold label, only when `repeats > 1`.
#'   \item `train_idx`    — list-column of training row indices.
#'   \item `validate_idx` — list-column of validation row indices.
#'   \item `train`, `validate` — list-columns of subsets, `data.table` or
#'         `data.frame` according to `out_type` (only when
#'         `materialize = TRUE`).
#' }
#'
#' @seealso [nest_cv()] for the nested `data.table` variant.
#'
#' @import data.table
#' @export
#' @examples
#' # Prepare example data: Convert first 3 columns of iris dataset to long format and split
#' dt_split <- w2l_split(data = iris, cols = 1:3)
#' # dt_split is now a list containing 3 data tables for Sepal.Length, Sepal.Width, and Petal.Length
#'
#' # Example 1: Single cross-validation (no repeats)
#' split_cv(
#'   data = dt_split,      # Input list of split data
#'   v = 3,                # Set 3-fold cross-validation
#'   repeats = 1,          # Perform cross-validation once (no repeats)
#'   seed = 123            # Reproducible folds
#' )
#' # Returns a list where each element contains:
#' # - id: fold labels (Fold1, Fold2, Fold3)
#' # - train_idx / validate_idx: row indices of each fold
#' # - train / validate: training and validation subsets
#'
#' # Example 2: Repeated cross-validation
#' split_cv(
#'   data = dt_split,      # Input list of split data
#'   v = 3,                # Set 3-fold cross-validation
#'   repeats = 2,          # Perform cross-validation twice
#'   seed = 123
#' )
#' # Returns a list where each element contains:
#' # - id: repeat labels (Repeat1, Repeat2)
#' # - id2: fold labels (Fold1, Fold2, Fold3)
#' # - train_idx / validate_idx, train / validate
#'
#' # Example 3: Stratified CV, indices only (memory friendly)
#' res <- split_cv(dt_split, v = 5, strata = "Species", seed = 1,
#'                 materialize = FALSE)
#' # Rebuild the training set of fold 1 of the first dataset when needed
#' head(dt_split[[1]][res[[1]]$train_idx[[1]], ])
split_cv <- function(data,
                     v           = 10L,
                     repeats     = 1L,
                     strata      = NULL,
                     breaks      = 4L,
                     pool        = 0.1,
                     seed        = NULL,
                     materialize = TRUE,
                     out_type    = "dt") {

  # -- 1. Input validation (once, before the loop) -----------------------------
  # A data.frame is also a list, so it must be rejected explicitly; otherwise
  # its columns would be treated as datasets.
  if (!is.list(data) || is.data.frame(data))
    stop("`data` must be a list of data.frames/data.tables. ",
         "To use a single data.frame, wrap it in list(). Received: ",
         paste(class(data), collapse = "/"), call. = FALSE)
  if (length(data) == 0L)
    stop("`data` must not be empty.", call. = FALSE)
  v       <- .check_count(v, "v", 2L)
  repeats <- .check_count(repeats, "repeats", 1L)
  breaks  <- .check_count(breaks, "breaks", 2L)
  if (!is.numeric(pool) || length(pool) != 1L || is.na(pool) || pool < 0 || pool >= 0.5)
    stop("`pool` must be a single number in [0, 0.5).", call. = FALSE)
  .check_flag(materialize, "materialize")
  .check_choice(out_type, c("dt", "df"), "out_type")

  is_valid <- vapply(data, is.data.frame, logical(1L))
  if (!all(is_valid)) {
    bad_idx <- which(!is_valid)
    stop("All elements of `data` must be data.frames or data.tables.\n",
         "  Invalid element(s) at index: ", paste(bad_idx, collapse = ", "),
         "\n  Types received: ",
         paste(vapply(data[bad_idx], function(x) class(x)[1L], character(1L)),
               collapse = ", "), call. = FALSE)
  }

  too_small <- vapply(data, nrow, integer(1L)) < v
  if (any(too_small))
    stop("Dataset(s) at index ", paste(which(too_small), collapse = ", "),
         " have fewer rows than `v` (", v, ").", call. = FALSE)

  if (!is.null(strata)) {
    if (!is.character(strata) || length(strata) != 1L || is.na(strata))
      stop("`strata` must be a single character string (column name).", call. = FALSE)
    strata_missing <- vapply(data, function(df) !strata %in% names(df), logical(1L))
    if (any(strata_missing)) {
      missing_idx <- which(strata_missing)
      stop("The `strata` column '", strata,
           "' is absent from dataset(s) at index: ",
           paste(missing_idx, collapse = ", "),
           if (!is.null(names(data)))
             paste0(" (name(s): '", paste(names(data)[missing_idx],
                                          collapse = "', '"), "')"),
           call. = FALSE)
    }
  }

  # -- 2. Fold assignment for every dataset -------------------------------------
  run <- function() {
    out <- lapply(data, function(x) {
      x <- if (data.table::is.data.table(x)) x else data.table::as.data.table(x)
      .vfold_table(x, v = v, repeats = repeats, strata = strata,
                   breaks = breaks, pool = pool, materialize = materialize,
                   out_type = out_type)
    })
    names(out) <- names(data)
    out
  }
  # A non-NULL seed makes the folds reproducible without disturbing the
  # caller's RNG stream.
  .with_seed(seed, run())
}