Skip to contents

Leakage-aware custom cross-validation

a5_custom_folds.Rmd
library(pipeML)
library(dplyr)
#> 
#> Attaching package: 'dplyr'
#> The following objects are masked from 'package:stats':
#> 
#>     filter, lag
#> The following objects are masked from 'package:base':
#> 
#>     intersect, setdiff, setequal, union
library(doParallel)  # also attaches foreach, used by the fold functions with tunable parameters
#> Loading required package: foreach
#> Loading required package: iterators
#> Loading required package: parallel

Leakage-aware cross-validation

A central design principle of pipeML is to prevent information leakage during model training and evaluation. In many machine learning workflows, feature engineering steps are applied to the full dataset before cross-validation, which can inadvertently introduce information from the test folds into the training process. This leads to overoptimistic performance estimates.

To address this, pipeML provides built-in support for custom fold construction through the fold_construction_fun argument in compute_features.training.ML(). This mechanism allows feature engineering and preprocessing steps to be recomputed independently within each cross-validation fold, ensuring that test samples never influence the training process.

This capability is a core component of the pipeML pipeline, enabling leakage-aware model development for datasets where features depend on the full sample structure.

Why custom fold construction is important

In many biological and high-dimensional datasets, features are not independent variables but are derived from the data itself. Examples include:

  • correlation-based clustering
  • dimensionality reduction (e.g., PCA)
  • gene set enrichment or pathway scoring
  • transcription factor activity inference
  • aggregation of features across samples

If these transformations are applied to the entire dataset before cross-validation, the test samples influence how the feature space is constructed. As a result, the model indirectly “sees” information from the test data during training.

By recomputing these steps within each fold, pipeML ensures that:

  • training data are used to define the feature space
  • test samples are projected onto the learned space without influencing it
  • model evaluation reflects true out-of-sample performance

This approach closely mimics how the model would behave when applied to completely unseen data, producing more realistic performance estimates.

Inside each fold, pipeML also removes near-constant and highly correlated features (|r| > 0.9) from the features built on the training part, and keeps the same features in the held-out part. The same is done on the features built on all training samples for the final model. Without a fold construction function, the features are the same in every fold, so this filter is applied once on all training samples before the cross-validation. Set preprocess = FALSE to skip it.

Step 1 - Define a base feature function

Structure of the Base Feature Function

The function is designed around two operational modes:

  • Training Mode (structure is NULL): The function learns the feature structure from the input dataset.
  • Projection Mode (structure provided): The function applies a previously learned structure to new data.

Why the function returns two objects?

  • features → always returned; the transformed representation of the current dataset (training or test).
  • structure → the learned structure; needed to project test data in future steps.

This ensures a leakage-aware workflow:

  • Training data defines the feature space.
  • Test data is projected without altering the learned structure.

Structure of base function

  • data: features as rows, samples as columns
  • structure: precomputed structure (e.g., clusters, components); NULL for training
  • …: additional arguments specific to the algorithm used
compute_features_modular <- function(data, structure = NULL, ...) {

  # TRAINING MODE
  if (is.null(structure)) {

    # -------------------- REPLACE THIS BLOCK --------------------
    structure <- learn_structure(data, ...) # user-defined function
    # -------------------- REPLACE THIS BLOCK --------------------

  }

  # PROJECT MODE

  # -------------------- REPLACE THIS BLOCK --------------------
  features <- project_data(data, structure, ...) # user-defined function
  # -------------------- REPLACE THIS BLOCK --------------------

  return(list(features = features, structure = structure))
}

Here we illustrate a correlation-based feature computation using Weighted Gene Co-expression Network Analysis (WGCNA): genes are grouped into co-expression modules on the training samples, and each module is summarized by the first principal component of its genes. This is just an example: in practice, you can use any feature computation that depends on multiple samples, such as clustering, PCA, among others.

library(WGCNA)
compute_features_modular <- function(counts, power = NULL, modules = NULL) {

  ## Just preprocessing (IGNORE)
  rownames(counts) <- gsub("-", ".", rownames(counts))
  datExpr <- t(counts)
  cor <- WGCNA::cor

  # TRAINING MODE
  if (is.null(modules)) {

    # -------------------- REPLACE THIS BLOCK --------------------
    net <- WGCNA::blockwiseModules(datExpr, power = power)
    modules <- net$colors
    names(modules) <- colnames(datExpr)
    # -------------------- REPLACE THIS BLOCK --------------------

  }

  # PROJECT MODE

  # -------------------- REPLACE THIS BLOCK --------------------
  module_features <- sapply(sort(unique(modules)), function(mod) {
    genes <- names(modules[modules == mod])
    pc <- prcomp(datExpr[, genes, drop = FALSE])
    pc$x[, 1]
  })
  # -------------------- REPLACE THIS BLOCK --------------------

  ## Just formatting (IGNORE)
  colnames(module_features) <- paste0("Module_", sort(unique(modules)))

  return(list(features = as.matrix(module_features), structure = modules))
}

Step 2 - Make the function suitable for pipeML

We then need to extend and give the correct format to this function to make it suitable for running across folds inside pipeML.

This template provides a modular framework to prepare cross-validation folds for pipeML in a leakage-aware way. It separates training vs projection:

  • Training mode: computes features and learns the data structure from the training folds
  • Projection mode: applies the learned structure to held-out folds without influencing it.

Users can easily adapt this template by replacing their previous compute_features_modular function.

Note.

In pipeML data corresponds to samples as rows and features as columns. If your compute_features_modular() needs features as rows, make sure to t() inside this function.

The parameters data, folds, and bestune are handled by pipeML automatically once fold_construction_fun is set. Do not change these parameter names or remove them.

  • data: the training features plus the outcome: a target column for classification, time and event columns for survival (see Custom folds for survival tasks).
  • folds: named list with the training rows of each fold.
  • bestune: NULL while running the folds; set by pipeML when building the features of the final model (see Why is bestune needed without tunable parameters?).
  • … : Additional parameters passed to your feature function

The function must:

  • in fold mode (bestune = NULL), save each fold as Results/fold_<fold name>.rds, with the training features (plus the outcome), the test features, the observed test outcome, the test row indices (rowIndex) and the fold name.
  • in final mode (bestune not NULL), return a list with the features of all training samples (plus the outcome), any output you want to keep (available as res$Custom_output) and bestune.

Make sure compute_features_modular() returns a matrix with the features. If not, make sure to extract them before adding the target column.

prepare_custom_folds <- function(data, folds = NULL, bestune = NULL, ...) {

  if (!is.null(bestune)) {

    obs_train <- data$target
    data$target <- NULL

    # -------------------- REPLACE THIS BLOCK --------------------
    result <- compute_features_modular(data, ...)
    # -------------------- REPLACE THIS BLOCK --------------------

    train_features_final = as.data.frame(result$features)
    train_features_final$target <- obs_train

    custom_output <- result

    return(list(train_features_final, custom_output, bestune))

  } else {

    processed_folds <- list()

    for (i in seq_along(folds)) {

      train_idx <- folds[[i]]
      test_idx  <- setdiff(seq_len(nrow(data)), train_idx)

      train_data <- data[train_idx, , drop = FALSE]
      obs_train <- train_data$target
      train_data$target <- NULL

      # -------------------- REPLACE THIS BLOCK --------------------
      train_result <- compute_features_modular(train_data, ...)
      # -------------------- REPLACE THIS BLOCK --------------------

      train_features = as.data.frame(train_result$features)
      train_features$target <- obs_train

      test_data <- data[test_idx, , drop = FALSE]
      obs_test <- test_data$target
      test_data$target <- NULL

      # -------------------- REPLACE THIS BLOCK --------------------
      test_features <- compute_features_modular(
        test_data,
        structure = train_result$structure,
        ...
      )
      # -------------------- REPLACE THIS BLOCK --------------------

      test_features = as.data.frame(test_features$features)

      processed_folds[[i]] <- list(
        train_data = train_features,
        test_data  = test_features,
        obs_test   = obs_test,
        rowIndex   = test_idx,
        fold_name  = names(folds)[i]
      )
    }

    for (i in seq_along(processed_folds)) {
      filename <- file.path("Results", paste0("fold_", names(folds)[i], ".rds"))
      saveRDS(processed_folds[[i]], file = filename)
    }

    return(processed_folds)
  }
}

Here we illustrate how the function will look applying our compute_features_modular() function.

Notice that each time I call the function compute_features_modular() I am setting my additional argument power:

prepare_WGCNA_folds <- function(data, folds = NULL, bestune = NULL, power) {

  if (!is.null(bestune)) {

    obs_train <- data$target
    data$target <- NULL

    # -------------------- REPLACE THIS BLOCK --------------------
    wgcna_result <- compute_features_modular(t(data), power = power)
    # -------------------- REPLACE THIS BLOCK --------------------

    train_cell_data_final <- as.data.frame(wgcna_result$features)
    train_cell_data_final$target <- obs_train

    custom_output <- wgcna_result

    return(list(train_cell_data_final, custom_output, bestune))

  } else {

    processed_folds <- list()

    for (i in seq_along(folds)) {

      train_idx <- folds[[i]]
      test_idx  <- setdiff(seq_len(nrow(data)), train_idx)

      train_data <- data[train_idx, , drop = FALSE]
      obs_train <- train_data$target
      train_data$target <- NULL

      # -------------------- REPLACE THIS BLOCK --------------------
      train_result <- compute_features_modular(t(train_data), power = power)
      # -------------------- REPLACE THIS BLOCK --------------------

      train_features <- as.data.frame(train_result$features)
      train_features$target <- obs_train

      test_data <- data[test_idx, , drop = FALSE]
      obs_test <- test_data$target
      test_data$target <- NULL

      # -------------------- REPLACE THIS BLOCK --------------------
      test_features <- compute_features_modular(t(test_data),
                                                modules = train_result$structure)
      # -------------------- REPLACE THIS BLOCK --------------------

      test_features <- as.data.frame(test_features$features)

      processed_folds[[i]] <- list(
        train_data = train_features,
        test_data  = test_features,
        obs_test   = obs_test,
        rowIndex   = test_idx,
        fold_name  = names(folds)[i]
      )
    }

    for (i in seq_along(processed_folds)) {
      filename <- file.path("Results", paste0("fold_", names(folds)[i], ".rds"))
      saveRDS(processed_folds[[i]], file = filename)
    }

    return(processed_folds)
  }
}

Load data example. counts_example contains genes as rows and samples as columns; we keep the samples annotated in coldata_example, in the same order:

coldata = pipeML::coldata_example
counts = pipeML::counts_example[, rownames(coldata)]

set.seed(123)

train_idx <- caret::createDataPartition(coldata$Response, p = 0.7, list = FALSE)

counts_train <- counts[, train_idx]
counts_test  <- counts[, -train_idx]

coldata_train <- coldata[train_idx, , drop = FALSE]
coldata_test  <- coldata[-train_idx, , drop = FALSE]

Step 3 - Run custom k-fold cross-validation

Once your custom fold function is ready, pass it to compute_features.training.ML() via fold_construction_fun.

The argument fold_construction_args_fixed corresponds to the additional parameters of your fold function (in this example power in prepare_WGCNA_folds()), set to the value to use. If your function does not have any additional parameters, you can omit this argument.

res_custom <- compute_features.training.ML(features_train = t(counts_train),
                                           target_var     = coldata_train$Response,
                                           task_type      = "classification",
                                           trait.positive = "R",
                                           metric         = "AUROC",
                                           k_folds        = 2,
                                           n_rep          = 1,
                                           return         = FALSE,
                                           fold_construction_fun        = prepare_WGCNA_folds,
                                           fold_construction_args_fixed = list(power = 6))

Notes on what pipeML does with a custom fold function:

  • fold files left in Results/ by a previous interrupted run are removed before your function is called, and the new ones are removed once they have been read;
  • for classification, the models are trained on the folds one after the other: the ncores argument of compute_features.training.ML() is not used. If building the features is slow, run it in parallel inside your fold function (see Tunable parameters within custom fold functions);
  • the hyperparameter grid of each model is the same in all folds. For random forest, the candidate values of mtry are sized from the fold with the fewest features built by your function.

Notice that res_custom$Custom_output contains the output of your base function on the full training set, in case it is needed (e.g. for prediction - see next step):

names(res_custom$Custom_output)
head(res_custom$Custom_output$features)
head(res_custom$Custom_output$structure)

Step 4 - Prediction on test data

To apply the model to new data, compute the same type of features, projecting the test samples on the structure learned from the training set (here, the WGCNA modules):

test = compute_features_modular(counts_test, modules = res_custom$Custom_output$structure)
test_features = test$features

Prediction

pred_custom <- compute_prediction(model = res_custom$Model,
                                  test_data = test_features,
                                  target_var = coldata_test$Response,
                                  task_type = "classification",
                                  trait.positive = "R",
                                  file.name = "Custom_fold")
pred_custom$AUC

Step 5 - SHAP values

SHAP values explain the final model, trained on the features computed on all training samples (res_custom$Custom_output$features), so each module has the same definition for all samples:

shap_custom <- compute_shap_values(model_trained = res_custom$Model,
                                   task_type = "classification")
head(shap_custom)

Tunable parameters within custom fold functions

In some scenarios, the feature construction step may include parameters whose values can influence model performance. For example, when computing WGCNA, parameters such as soft-thresholding power, minimum module size, module merging threshold, and module splitting sensitivity may affect the resulting features and therefore the downstream model performance.

To address this, pipeML supports hyperparameter tuning within your custom fold functions. This allows users to identify which parameter values lead to the best predictive performance.

Briefly, the user provides a grid of candidate parameter values, and pipeML will:

  • construct custom cross-validation folds for each parameter combination
  • train and evaluate machine learning models for each configuration
  • compare the resulting performance across folds and repetitions
  • return the parameter values that maximize the selected performance metric

Before that, user needs to modify compute_features_modular and prepare_custom_folds() to account for all these parameters.

For the compute_features_modular we are only going to add the tunable parameters in our function call:

compute_features_modular <- function(
    counts,
    power = NULL,
    modules = NULL,
    ## tunable parameters
    minModuleSize = 20,
    mergeCutHeight = 0.15,
    deepSplit = 2
) {

  ## Just preprocessing (IGNORE)
  rownames(counts) <- gsub("-", ".", rownames(counts))
  datExpr <- t(counts)
  cor <- WGCNA::cor

  # TRAINING MODE

  if (is.null(modules)) {

    # -------------------- REPLACE THIS BLOCK --------------------
    net <- WGCNA::blockwiseModules(
      datExpr,
      power = power,
      minModuleSize = minModuleSize,
      mergeCutHeight = mergeCutHeight,
      deepSplit = deepSplit
    )
    modules <- net$colors
    names(modules) <- colnames(datExpr)
    # -------------------- REPLACE THIS BLOCK --------------------

  }

  # PROJECT MODE

  # -------------------- REPLACE THIS BLOCK --------------------
  module_features <- sapply(sort(unique(modules)), function(mod) {
    genes <- names(modules[modules == mod])
    pc <- prcomp(datExpr[, genes, drop = FALSE])
    pc$x[, 1]
  })
  # -------------------- REPLACE THIS BLOCK --------------------

  ## Just formatting (IGNORE)
  colnames(module_features) <- paste0("Module_", sort(unique(modules)))

  return(list(features = as.matrix(module_features), structure = modules))
}

Then we will use this version of prepare_custom_folds(), modified to account for the parameter combinations. Compared to the version without tunable parameters:

  • in fold mode, each Results/fold_<fold name>.rds file contains a list with one element per parameter combination, and each element also stores the combination in params.
  • in final mode, the selected values are read from bestune, and the function returns them (as a data frame) as the third element.

Notice we have an additional parameter ncores that controls parallelization when evaluating all runs (do not remove it!).

prepare_custom_folds_tuning <- function(data,
                                        folds = NULL,
                                        bestune = NULL,
                                        ncores = NULL, ...){

  if (!is.null(bestune)) {

    obs_train <- data$target
    data$target <- NULL

    # -------------------- REPLACE THIS BLOCK --------------------
    required_cols <- c() # list your params separated by a comma
    # -------------------- REPLACE THIS BLOCK --------------------

    best_params <- if (is.data.frame(bestune)) {
      if (all(required_cols %in% names(bestune))) {
        dplyr::select(bestune, dplyr::all_of(required_cols))
      }else{stop("Not all tunable params found. Verify your function.")}
    } else if (is.list(bestune)) {
      if (all(required_cols %in% names(bestune))) {
        tibble::as_tibble(bestune[required_cols])
      }else{stop("Not all tunable params found. Verify your function.")}
    } else {
      stop("`bestune` must be a data.frame or list.")
    }

    # -------------------- REPLACE THIS BLOCK --------------------
    res_final <- compute_features_modular(
      counts = data,
      # change param1, param2, ... for the names of your parameters
      param1 = best_params$param1,
      param2 = best_params$param2,
      param3 = best_params$param3,
      ...
    )
    # -------------------- REPLACE THIS BLOCK --------------------

    train_features_final = as.data.frame(res_final$features)
    train_features_final$target <- obs_train

    custom_output <- res_final

    return(list(train_features_final, custom_output, best_params))

  } else {

    # -------------------- REPLACE THIS BLOCK --------------------
    custom_grid <- expand.grid(
      param1 = param1,
      param2 = param2,
      param3 = param3,
      ...,
      stringsAsFactors = FALSE
    )
    # -------------------- REPLACE THIS BLOCK --------------------

    if (is.null(ncores)) ncores <- parallel::detectCores() - 2
    cl <- parallel::makeCluster(ncores)
    doParallel::registerDoParallel(cl)

    processed_folds <- foreach::foreach(i = seq_along(folds),
                                        .packages = c("dplyr"),
                                        .export = c("compute_features_modular")
                                       ) %dopar% {

      train_idx <- folds[[i]]
      test_idx  <- setdiff(seq_len(nrow(data)), train_idx)

      train_data <- data[train_idx, , drop = FALSE]
      obs_train <- train_data$target
      train_data$target <- NULL

      fold_results <- lapply(seq_len(nrow(custom_grid)), function(j) {
        params <- custom_grid[j, , drop = FALSE]

        # -------------------- REPLACE THIS BLOCK --------------------
        res_train <- compute_features_modular(
          counts = train_data,
          param1 = params$param1,
          param2 = params$param2,
          param3 = params$param3,
          ...
        )
        # -------------------- REPLACE THIS BLOCK --------------------

        train_features <- as.data.frame(res_train$features)
        train_features$target <- obs_train

        test_data <- data[test_idx, , drop = FALSE]
        obs_test <- test_data$target
        test_data$target <- NULL

        # -------------------- REPLACE THIS BLOCK --------------------
        test_features <- compute_features_modular(
          test_data,
          structure = res_train$structure,
          ...
        )
        # -------------------- REPLACE THIS BLOCK --------------------

        test_features = as.data.frame(test_features$features)

        list(
          train_data = train_features,
          test_data  = test_features,
          obs_test   = obs_test,
          rowIndex   = test_idx,
          fold_name  = names(folds)[i],
          params     = params
        )
      })

      filename <- file.path("Results", paste0("fold_", names(folds)[i], ".rds"))
      saveRDS(fold_results, file = filename)

      fold_results
    }

    parallel::stopCluster(cl)
    foreach::registerDoSEQ()
    gc()

  }
}

In our case it would be:

prepare_WGCNA_folds_modular <- function(
    data,
    folds = NULL,
    bestune = NULL,
    power = NULL,
    ncores = NULL,
    ### tunable parameters
    minModuleSize,
    mergeCutHeight,
    deepSplit
) {

  if (!is.null(bestune)) {

    obs_train <- data$target
    data$target <- NULL

    # -------------------- REPLACE THIS BLOCK --------------------
    required_cols <- c("minModuleSize", "mergeCutHeight", "deepSplit")
    # -------------------- REPLACE THIS BLOCK --------------------

    best_params <- if (is.data.frame(bestune)) {
      if (all(required_cols %in% names(bestune))) {
        dplyr::select(bestune, dplyr::all_of(required_cols))
      }else{stop("Not all tunable params found. Verify your function.")}
    } else if (is.list(bestune)) {
      if (all(required_cols %in% names(bestune))) {
        tibble::as_tibble(bestune[required_cols])
      }else{stop("Not all tunable params found. Verify your function.")}
    } else {
      stop("`bestune` must be a data.frame or list.")
    }

    # -------------------- REPLACE THIS BLOCK --------------------
    res_final <- compute_features_modular(
      counts = t(data),
      power = power,
      ## tunable parameters
      minModuleSize = best_params$minModuleSize,
      mergeCutHeight = best_params$mergeCutHeight,
      deepSplit = best_params$deepSplit
    )
    # -------------------- REPLACE THIS BLOCK --------------------

    train_cell_data_final <- as.data.frame(res_final$features)
    train_cell_data_final$target <- obs_train

    custom_output <- res_final

    return(list(train_cell_data_final, custom_output, best_params))

  } else {

    # -------------------- REPLACE THIS BLOCK --------------------
    custom_grid <- expand.grid(
      minModuleSize = minModuleSize,
      mergeCutHeight = mergeCutHeight,
      deepSplit = deepSplit,
      stringsAsFactors = FALSE
    )
    # -------------------- REPLACE THIS BLOCK --------------------

    if (is.null(ncores)) ncores <- parallel::detectCores() - 2
    cl <- parallel::makeCluster(ncores)
    doParallel::registerDoParallel(cl)

    processed_folds <- foreach::foreach(i = seq_along(folds),
                                        .packages = c("dplyr"),
                                        .export = c("compute_features_modular")
                                        ) %dopar% {

      train_idx <- folds[[i]]
      test_idx <- setdiff(seq_len(nrow(data)), train_idx)

      train_data <- data[train_idx, , drop = FALSE]
      obs_train <- train_data$target
      train_data$target <- NULL

      fold_results <- lapply(seq_len(nrow(custom_grid)), function(j) {

        params <- custom_grid[j, , drop = FALSE]

        # -------------------- REPLACE THIS BLOCK --------------------
        wgcna_train <- compute_features_modular(
          counts = t(train_data),
          power = power,
          ## tunable parameters
          minModuleSize = params$minModuleSize,
          mergeCutHeight = params$mergeCutHeight,
          deepSplit = params$deepSplit
        )
        # -------------------- REPLACE THIS BLOCK --------------------

        train_features <- as.data.frame(wgcna_train$features)
        train_features$target <- obs_train

        test_data <- data[test_idx, , drop = FALSE]
        obs_test <- test_data$target
        test_data$target <- NULL

        # -------------------- REPLACE THIS BLOCK --------------------
        wgcna_test <- compute_features_modular(
          counts = t(test_data),
          modules = wgcna_train$structure
        )
        # -------------------- REPLACE THIS BLOCK --------------------

        test_features <- as.data.frame(wgcna_test$features)

        list(
          train_data = train_features,
          test_data = test_features,
          obs_test = obs_test,
          rowIndex = test_idx,
          fold_name = names(folds)[i],
          params = params
        )
      })

      filename <- file.path("Results", paste0("fold_", names(folds)[i], ".rds"))
      saveRDS(fold_results, file = filename)

      fold_results
    }

    parallel::stopCluster(cl)
    foreach::registerDoSEQ()
    gc()

  }
}

To enable this functionality, the user must specify the argument fold_construction_args_tunable, which contains the set of parameter values to be evaluated.

The argument fold_construction_args_fixed corresponds to the additional parameters of the fold function that are not tunable (in our case power and ncores). These parameters have the same value across all runs (fixed). Be sure of setting ncores to a value that your computer can handle to avoid crashing.

The tunable parameters define the search space explored during feature computation. The total number of configurations corresponds to all combinations of the provided parameter values. In this example:

  • minModuleSize = c(20, 50) → 2 values
  • mergeCutHeight = 0.25 → 1 value
  • deepSplit = c(1, 2) → 2 values

This results in 2 × 1 × 2 = 4 feature parameter combinations.

For each of these configurations, the machine learning models are trained and tuned. For example, for logistic regression with elastic net (glmnet), pipeML evaluates 2 alpha values (0 and 1) and 20 lambda values, which results in 40 model configurations for each feature combination, so 4 × 40 = 160 model fits, further multiplied by the number of folds and repetitions. The same happens for each of the other algorithms, so the running time grows quickly with the number of feature parameter combinations.

res_params <- compute_features.training.ML(features_train = t(counts_train),
                                           target_var     = coldata_train$Response,
                                           task_type = "classification",
                                           trait.positive = "R",
                                           metric = "AUROC",
                                           k_folds = 2,
                                           n_rep = 1,
                                           return = FALSE,
                                           fold_construction_fun = prepare_WGCNA_folds_modular,
                                           fold_construction_args_fixed = list(power = 6,
                                                                               ncores = 2),
                                           fold_construction_args_tunable = list(
                                             minModuleSize = c(20, 50),
                                             mergeCutHeight = 0.25,
                                             deepSplit = c(1, 2)
                                           ))

pipeML will automatically train the model with the combination of parameters which maximizes the metric chosen. The selected parameters are in res_params$Custom_output$Parameters:

res_params$Custom_output$Parameters

Prediction on the test set works as without tunable parameters, projecting the test samples on the modules learned with the selected parameters:

test_params <- compute_features_modular(counts_test, modules = res_params$Custom_output$structure)

pred_params <- compute_prediction(model = res_params$Model,
                                  test_data = test_params$features,
                                  target_var = coldata_test$Response,
                                  task_type = "classification",
                                  trait.positive = "R")
pred_params$AUC

SHAP values explain the final model, built with the selected parameters:

shap_params <- compute_shap_values(model_trained = res_params$Model,
                                   task_type = "classification")
head(shap_params)

Custom folds for survival tasks

Custom fold functions work the same way for survival tasks. The only difference is the outcome: data contains the survival time and event in two columns named time and event (instead of target). As for classification, the function must:

  • remove time and event from data before building the features;
  • add them back to the training and test features of each fold, and to the training features returned in final mode.

pipeML stops with an error if a feature returned by the function is a copy of time or event (e.g. because they were not removed before building the features).

Here we build features with a principal component analysis (PCA) learned on the training samples, on the survival example data. First, the base feature function:

compute_pca_features <- function(data, n_comp = 2, structure = NULL) {

  # TRAINING MODE: learn the principal components
  if (is.null(structure)) {
    structure <- prcomp(data, center = TRUE, scale. = TRUE)
  }

  # PROJECT MODE: project the samples on the learned components
  features <- predict(structure, newdata = data)[, seq_len(n_comp), drop = FALSE]

  return(list(features = as.data.frame(features), structure = structure))
}

Data:

data = pipeML::data_example_survival
X <- data %>% dplyr::select(-time, -status)
time <- data$time
event <- data$status

set.seed(123)
train_idx <- caret::createDataPartition(event, p = 0.7, list = FALSE)

X_train <- X[train_idx, ]
X_test  <- X[-train_idx, ]
time_train <- time[train_idx]
time_test  <- time[-train_idx]
event_train <- event[train_idx]
event_test  <- event[-train_idx]

Without tunable parameters. The number of components is fixed, and passed as a fixed argument:

prepare_PCA_folds <- function(data, folds = NULL, bestune = NULL, n_comp = 2) {

  # Outcome, and features only
  time <- data$time
  event <- data$event
  data <- data[, setdiff(colnames(data), c("time", "event")), drop = FALSE]

  if (!is.null(bestune)) {
    # Final mode: features of all training samples
    result <- compute_pca_features(data, n_comp = n_comp)
    train_features <- result$features
    train_features$time <- time
    train_features$event <- event
    return(list(train_features, result, bestune))
  }

  # Fold mode
  for (i in seq_along(folds)) {
    train_idx <- folds[[i]]
    test_idx  <- setdiff(seq_len(nrow(data)), train_idx)

    train_result <- compute_pca_features(data[train_idx, , drop = FALSE], n_comp = n_comp)
    train_features <- train_result$features
    train_features$time <- time[train_idx]
    train_features$event <- event[train_idx]

    test_features <- compute_pca_features(data[test_idx, , drop = FALSE], n_comp = n_comp,
                                          structure = train_result$structure)$features
    test_features$time <- time[test_idx]
    test_features$event <- event[test_idx]

    saveRDS(list(train_data = train_features, test_data = test_features, rowIndex = test_idx),
            file.path("Results", paste0("fold_", names(folds)[i], ".rds")))
  }
}
res_pca <- compute_features.training.ML(features_train = X_train,
                                        task_type = "survival",
                                        time_var = time_train,
                                        event_var = event_train,
                                        k_folds = 2,
                                        n_rep = 1,
                                        ncores = 2,
                                        fold_construction_fun = prepare_PCA_folds,
                                        fold_construction_args_fixed = list(n_comp = 5))

Prediction: project the test samples on the components learned on the training set, then predict:

test_pca <- compute_pca_features(X_test, n_comp = 5, structure = res_pca$Custom_output$structure)

pred_pca <- compute_prediction(model = res_pca$Model,
                               test_data = test_pca$features,
                               task_type = "survival",
                               time_var = time_test,
                               event_var = event_test)
pred_pca$c_index

SHAP values:

shap_pca <- compute_shap_values(model_trained = res_pca$Model,
                                task_type = "survival")
head(shap_pca)

With tunable parameters. Here the number of components is tuned. As for classification, each fold file contains one element per candidate value (with the value in params), and in final mode the selected value is read from bestune and returned as the third element:

prepare_PCA_folds_tuning <- function(data, folds = NULL, bestune = NULL, n_comp) {

  # Outcome, and features only
  time <- data$time
  event <- data$event
  data <- data[, setdiff(colnames(data), c("time", "event")), drop = FALSE]

  if (!is.null(bestune)) {
    # Final mode: features of all training samples with the selected number of components
    best_n_comp <- bestune$n_comp[1]
    result <- compute_pca_features(data, n_comp = best_n_comp)
    train_features <- result$features
    train_features$time <- time
    train_features$event <- event
    return(list(train_features, result, data.frame(n_comp = best_n_comp)))
  }

  # Fold mode: one element per candidate number of components
  for (i in seq_along(folds)) {
    train_idx <- folds[[i]]
    test_idx  <- setdiff(seq_len(nrow(data)), train_idx)

    fold_results <- lapply(n_comp, function(k) {
      train_result <- compute_pca_features(data[train_idx, , drop = FALSE], n_comp = k)
      train_features <- train_result$features
      train_features$time <- time[train_idx]
      train_features$event <- event[train_idx]

      test_features <- compute_pca_features(data[test_idx, , drop = FALSE], n_comp = k,
                                            structure = train_result$structure)$features
      test_features$time <- time[test_idx]
      test_features$event <- event[test_idx]

      list(train_data = train_features, test_data = test_features, rowIndex = test_idx,
           params = data.frame(n_comp = k))
    })

    saveRDS(fold_results, file.path("Results", paste0("fold_", names(folds)[i], ".rds")))
  }
}
res_pca_tuning <- compute_features.training.ML(features_train = X_train,
                                               task_type = "survival",
                                               time_var = time_train,
                                               event_var = event_train,
                                               k_folds = 2,
                                               n_rep = 1,
                                               ncores = 2,
                                               fold_construction_fun = prepare_PCA_folds_tuning,
                                               fold_construction_args_tunable = list(n_comp = c(3, 5)))

res_pca_tuning$Custom_output$Parameters   # selected number of components

Prediction with the selected number of components, and SHAP values:

best_n_comp <- res_pca_tuning$Custom_output$Parameters$n_comp
test_pca_tuning <- compute_pca_features(X_test, n_comp = best_n_comp,
                                        structure = res_pca_tuning$Custom_output$structure)

pred_pca_tuning <- compute_prediction(model = res_pca_tuning$Model,
                                      test_data = test_pca_tuning$features,
                                      task_type = "survival",
                                      time_var = time_test,
                                      event_var = event_test)
pred_pca_tuning$c_index

shap_pca_tuning <- compute_shap_values(model_trained = res_pca_tuning$Model,
                                       task_type = "survival")
head(shap_pca_tuning)

Why is bestune needed without tunable parameters?

For this to work, your custom function must accept a bestune argument, which is used internally to inject the optimized parameter values found during the tuning step.

When bestune is NULL, the function assumes that tuning has not yet been performed. A grid of candidate parameter values (defined in fold_construction_args_tunable) is generated. For each fold, the function iterates through all combinations of parameter values and recomputes the features. This exploration step can be parallelized across folds using foreach and doParallel, allowing multiple folds to be processed simultaneously. Parallelization reduces runtime considerably when the parameter grid or number of folds is large.

When bestune is not NULL, it means the tuning process has already been completed. The optimized parameter values are extracted from the bestune object. Features are then recomputed once on the full training dataset using these tuned parameters. This ensures the final model is trained with the best parameter setting identified during cross-validation.

In summary, the bestune argument acts as a control switch:

  • NULL → build the features of each cross-validation fold (for each parameter combination, if any).
  • non-NULL → lock in the tuned parameter values and rebuild the features for final training.

This design allows a single custom fold-construction function to handle both hyperparameter tuning (exploration, parallelized) and final model preparation (exploitation, single optimized run). Without tunable parameters, bestune is still used as this switch between fold mode and final mode.