"grouped" vfolding in rsample

group-level-vfolding.R

library(dplyr)
library(purrr)
library(rsample)

# suppose we want to keep cylinder-groups together
# we'll vfold those instead of the whole thing
initial_fold <- mtcars %>%
  distinct(cyl) %>%
  vfold_cv(v = 3)

# take a look
initial_fold %>%
  pull(splits) %>%
  map(assessment)

# take an existing rsplit object and expand given the original it was subsampled from
expand_split <- function(split, orig) {
  
  vars <- colnames(split$data)
  
  rows_in  <- which(pull(orig, vars) %in% pull(analysis(split), vars))
  rows_out <- which(pull(orig, vars) %in% pull(assessment(split), vars))
  
  rsample:::rsplit(orig, sample(rows_in), sample(rows_out))  # can't forget to shuffle
}

# now apply that to each split
expanded_fold <- initial_fold %>%
  mutate(splits = map(splits, expand_split, mtcars))

# take a look
expanded_fold %>%
  pull(splits) %>%
  map(assessment)

# other side
expanded_fold %>%
  pull(splits) %>%
  map(analysis)