Subgroup-specific cross-validation

NAMESPACE

@@ -1,12 +1,15 @@
 # Generated by roxygen2: do not edit by hand
 
+export("%notin%")
 export(combine_results)
 export(combiner_array)
 export(combiner_c)
 export(combiner_factor)
 export(combiner_rbind)
 export(cross_validate)
 export(fold_from_foldvec)
+export(fold_from_foldvec_subgroup)
+export(fold_from_foldvec_subgroup_stratoutcome)
 export(fold_index)
 export(folds2foldvec)
 export(folds_bootstrap)
@@ -17,6 +20,7 @@ export(folds_rolling_origin)
 export(folds_rolling_origin_pooled)
 export(folds_rolling_window)
 export(folds_rolling_window_pooled)
+export(folds_subgroup)
 export(folds_vfold)
 export(folds_vfold_rolling_origin_pooled)
 export(folds_vfold_rolling_window_pooled)

R/fold_funs.R

@@ -29,6 +29,8 @@
 #'  subject in the sample.
 #' @param id An optional vector of unique identifiers corresponding to the time
 #'  vector. These can be used to subset the time vector.
+#'  @param subgroup Binary indicator vector of membership in the subgroup of interest for subgroup-specific cross-validation (eg data$S).
+#'  @param strat_outcome Optional binary outcome vector (eg data$Y) for subgroup-specific cross-validation if equal proportion of outcome in training and validation folds desired.
 #'
 #' @importFrom assertthat assert_that
 #'
@@ -369,6 +371,50 @@ folds_vfold_rolling_window_pooled <- function(n, t, id = NULL, time = NULL,
   return(folds)
 }
 
+#########################################################
+#Add option to make subgroup-specific cross-validation folds
+
+#' @rdname fold_funs
+#' @export
+folds_subgroup <- function(n, V = 10L, subgroup, strat_outcome = NULL) {
+
+  if(is.null(strat_outcome)==TRUE){
+
+    subgroup_index <- which(subgroup == 1)
+
+    folds <- rep(seq_len(V), length = length(subgroup_index))
+
+    # shuffle folds
+    folds <- sample(folds)
+    foldsdf <- data.frame(subgroup_index, folds)
+
+    # generate fold vectors
+    folds <- lapply(seq_len(V), fold_from_foldvec_subgroup, foldsdf, n)
+    return(folds)
+
+  }
+  else {
+    subgroup_index1 <- which(subgroup == 1 & strat_outcome == 1)
+    subgroup_index2 <- which(subgroup == 1 & strat_outcome == 0)
+
+    #Two sets of folds, one for outcome == 1 and one for outcome == 0
+    folds1 <- rep(seq_len(V), length = length(subgroup_index1))
+    folds2 <- rep(seq_len(V), length = length(subgroup_index2))
+
+    # shuffle folds
+    folds1 <- sample(folds1)
+    foldsdf1 <- data.frame(subgroup_index1, folds1)
+
+    folds2 <- sample(folds2)
+    foldsdf2 <- data.frame(subgroup_index2, folds2)
+
+    # generate fold vectors
+    folds <- lapply(seq_len(V), fold_from_foldvec_subgroup_stratoutcome, foldsdf1, foldsdf2,  n)
+    return(folds)
+  }
+}
+
+
 ###############################################################################
 
 #' Check ID and Time Compatibility
@@ -411,3 +457,4 @@ check_id_and_time <- function(id, time) {
     )
   }
 }
+

R/utils.R

@@ -64,3 +64,22 @@ folds2foldvec <- function(folds) {
   fold_id <- nums[order(fold_index)]
   return(fold_id)
 }
+
+################################################################################
+
+#' @export
+`%notin%` <- Negate(`%in%`)
+
+#' @export
+fold_from_foldvec_subgroup <- function(v, foldsdf, n) {
+  validation_set <- foldsdf[,1][which(foldsdf[,2] == v)]
+  training_set <- which(seq_len(n) %notin% validation_set)
+  make_fold(v, training_set, validation_set)
+}
+
+#' @export
+fold_from_foldvec_subgroup_stratoutcome <- function(v, foldsdf1, foldsdf2, n) {
+  validation_set <- c(foldsdf1[,1][which(foldsdf1[,2] == v)], foldsdf2[,1][which(foldsdf2[,2] == v)])
+  training_set <- which(seq_len(n) %notin% validation_set)
+  make_fold(v, training_set, validation_set)
+}