R/build_fwe_correction.R

In UNCDEPENdLab/fmri.pipeline: Analysis of fMRI Data on HPC clusters

#' helper function to create an FWE object that specifies the type of FWE correction and the 
#'   model outputs to which it is applied
#' @param gpa a \code{glm_pipeline_arguments} object that already contains model output information in $l3_model_setup$fsl
#' @param level an integer indicating what modeling level this FWE correction applies to (1, 2, or 3)
#' @param lg a logger object for logging messages
#' @details note that this only works for level = 3 for now
create_fwe_spec <- function(gpa, level = level, lg = NULL) {
  checkmate::assert_class(gpa, "glm_pipeline_arguments")
  checkmate::assert_integerish(level, lower = 1, upper = 3)
  checkmate::assert_class(lg, "Logger")

  cat("Welcome to the FWE setup menu\n")
  cat("At present, we offer 4 forms of FWE correction: pTFCE, AFNI 3dFWHMx + 3dClustSim, 3dttest++ residual permutation + 3dClustSim, and PALM permutation\n")

  which_fwe <- select.list(
    choices = c("pTFCE", "3dFWHMx + 3dClustSim", "3dtest++ -randomsign + 3dClustSim", "PALM"),
    title = c("Which correction would you like to add?")
  )

  stopifnot(level == 3L) # only support 3rd-level FWE for now

  # need to choose to which models this applies...
  # for now, skip the model lookup from the run step and just rely on l3_model_setup
  # if (level == 3L) {
  #   l3_set <- choose_glm_models(gpa, "prompt", level = 3L, lg = lg)
  #   if (is.null(l3_set)) {
  #     message("No l3 models were selected for this FWE.")
  #     return(NULL) #maybe it's okay not to have any models?
  #   }
  # }

  fwe_set <- list()
  fields <- c("l1_model", "l1_cope_name", "l2_model", "l2_cope_name", "l3_model")
  for (ff in fields) {
    fwe_set[[ff]] <- select.list(unique(gpa$l3_model_setup$fsl[[ff]]),
      multiple = TRUE,
      title = paste("For this FWE, choose which", ff, "to include in this FWE.")
    )
  }

  # build filter statement
  filter_expr <- paste(lapply(seq_along(fwe_set), function(x) {
    paste(names(fwe_set)[x], "%in%", paste(deparse(fwe_set[[x]]), collapse = ""))
  }), collapse = " & ")

  # https://www.r-bloggers.com/2020/05/filtering-with-string-statements-in-dplyr/
  to_fwe <- gpa$l3_model_setup$fsl %>% dplyr::filter(rlang::eval_tidy(rlang::parse_expr(filter_expr)))

  if (nrow(to_fwe) == 0L) {
    msg <- "No model outputs match this combination of selections."
    lg$error(msg)
    stop(msg)
  }

  if (any(to_fwe$feat_complete == FALSE)) {
    lg$warn("Some L3 models requested are currently incompleted, preventing FWE correction.")
    lg$warn("%s", capture.output(print(to_fwe %>% filter(feat_complete == FALSE))))
    to_fwe <- to_fwe %>% dplyr::filter(feat_complete == TRUE)
    if (nrow(to_fwe) == 0L) {
      msg <-"All L3 models requested are incomplete and this function cannot continue."
      lg$error(msg)
      stop(msg)
    }
  }


  # I guess we could have lower level subsetting, too.
  # gfeat_set <- gpa$l3_model_setup$fsl %>% dplyr::filter(l3_model == !!l3_set & feat_complete == TRUE)
  
  #choices = c("pTFCE", "3dFWHMx + 3dClustSim", "3dtest++ -randomsign + 3dClustSim", "PALM"),

  # all objects here should support $submit(force = X)
  if (which_fwe == "pTFCE") {
    fwe_obj <- ptfce_spec$new(gfeat_dir = to_fwe$feat_dir)
  } else if (which_fwe == "3dFWHMx + 3dClustSim") {
    warning("3dFWHMx is not quite in here yet!")
    fwe_obj <- NULL
  } else if (which_fwe == "3dtest++ -randomsign + 3dClustSim") {
    fwe_obj <- build_3dclustsim_permutation(to_fwe)
  } else if (which == "PALM") {
    warning("PALM is not quite in here yet!")
    fwe_obj <- NULL
  }

  return(fwe_obj)

}

#' helper function for setting up 3dClustSim objects for permutation (via 3dttest++).
#' @param to_fwe a data.frame containing model information for the group analyses to be passed through 3dClustSim FWE
#' @keywords internal
build_3dclustsim_permutation <- function(to_fwe) {
    ## NUMBER OF TOTAL ITERATIONS FOR 3dClustSim
    iter <- NULL
    cat(
      "",
      "The permutation-based approach to 3dClustSim uses null datasets generated by sign-flipping and/or permutation",
      "  of the residuals from the group analysis. In FSL, these residuals are provided in the res4d.nii.gz inside the",
      "  cope<X>.feat subfolder of the group .gfeat folder. For 3dClustSim, how many null datasets should be generated from",
      "  the group residuals? The default is 30000, but higher numbers (e.g., 100000) are better if you want to test lower",
      "  voxelwise thresholds.",
      "",
      sep = "\n"
    )
    while (is.null(iter)) {
      prompt <- "How many null datasets should be generated for 3dClustSim?\n(Press ENTER to accept the default of 30000)?\n"
      ii <- readline(prompt)
      if (ii == "") {
        iter <- 30000
      } else {
        ii <- as.integer(ii)
        if (!checkmate::test_integerish(ii, lower = 10, upper = 1e8)) {
          cat("Enter a number between 10 and 10 million\n")
        } else {
          iter <- ii
        }
      }
    }

    ## NUMBER OF JOBS FOR DISTRIBUTION NULL DATASET CALCULATIONS
    cat(
      glue(
        "",
        "Null dataset generation can be split across jobs on the cluster to reduce the calculation time.",
        "The total number of datasets is {as.integer(iter)}. These can be divided across jobs. For example, if distributed across",
        "5 jobs, the number iterations per job would be approximately {ceiling(iter/5)}.",
        "",
        .sep = "\n"
      )
    )

    residual_njobs <- NULL
    while (is.null(residual_njobs)) {
      prompt <- "How many jobs should be used for creating the sign-flipping/permutation datasets?\n(Press ENTER to accept the default of 32)?\n"
      ii <- readline(prompt)
      if (ii == "") {
        residual_njobs <- 32L
      } else {
        ii <- as.integer(ii)
        if (!checkmate::test_integerish(ii, lower = 1, upper = 1e5)) {
          cat("Enter a number between 1 and 100,000\n")
        } else {
          residual_njobs <- ii
        }
      }
    }

    ## NUMBER OF THREADS FOR 3dClustSim ITSELF
    cat(
      "",
      "Clusterizing and thresholding null datasets within 3dClustSim can be expedited by distributing the computations",
      "  across multiple compute threads (think: cores on the CPU). Note that this is controlled by the OMP_NUM_THREADS",
      "  environment variable. Here, the number chosen also controls how many cores are requested by a job on the scheduler.",
      "",
      sep="\n"
    )

    ncpus <- NULL
    while (is.null(ncpus)) {
      prompt <- "How many cores should be used for running 3dClustSim?\n(Press ENTER to accept the default of 8)?\n"
      ii <- readline(prompt)
      if (ii == "") {
        ncpus <- 8
      } else {
        ii <- as.integer(ii)
        if (!checkmate::test_integerish(ii, lower = 1, upper = 1e3)) {
          cat("Enter a number between 1 and 1,000\n")
        } else {
          ncpus <- ii
        }
      }
    }

    cat(
      "For 3dClustSim results to be meaningful, it is important to specify the volume over which one wishes to correct the FWE.",
      "  Said differently, we want to use a mask containing the voxels of interest and only correct over those voxelse, rather",
      "  than the entire space of the image (including many empty voxels outside the brain).", 
      "  Here, you can specify the path to a specific mask over which you want to correct the FWE. If you do not choose a mask",
      "  the default mask.nii.gz file in the .feat folder for your analysis will be used.",
      sep="\n"
    )

    clustsim_mask <- NULL
    while (is.null(clustsim_mask)) {
      prompt <- "Specify the path to a mask file (NIfTI) to be used for 3dClustSim correction.\n(Press ENTER to accept the default FSL mask)\n"
      ii <- readline(prompt)
      if (ii == "") {
        clustsim_mask <- "..DEFAULT.."
      } else {
        if (!checkmate::test_file_exists(ii)) {
          cat(glue("Cannot find specified mask file: {ii}\n", .trim = FALSE))
        } else {
          # always use full path
          clustsim_mask <- normalizePath(ii)
        }
      }
    }

    cat(
      "",
      "3dClustSim can calculate FWE-corrected cluster thresholds across a range of voxelwise p-values.",
      "  Note that this is not an arbitrary choice and the use of liberal threshold (e.g., p < .05) has been",
      "  associated with poor/anticonservative FWE correction (see, for example, Woo, Krishnan, & Wager, 2014).",
      "  See 3dClustSim -pthr for documentation.",
      "",
      sep="\n"
    )
    pthr <- NULL    
    while (is.null(pthr)) {
      prompt <- "Specify the voxelwise p-values to be used in FWE cluster correction.\n(Press ENTER to accept default of .01 .005 .002 .001 .0005 .0002 .0001)\n"
      ii <- readline(prompt)
      if (ii == "") {
        pthr <- ".01 .005 .002 .001 .0005 .0002 .0001"
      } else {
        # validate input
        pthr <- check_nums(pthr, lower = 1e-10, upper = .999)
      }
    }

    cat(
      "",
      "3dClustSim can also calculate the number of voxels needed for an FWE-significant cluster across a range of cluster p-value thresholds.",
      "  This calculation indicates how many voxels are needed in a cluster at an FWE-corrected",
      "  whole-brain significance level of .05, .02, .01, and so on. See 3dClustSim -athr for documentation.",
      "",
      sep = "\n"
    )
    athr <- NULL
    while (is.null(athr)) {
      prompt <- "Specify the whole-brain p-values (alpha) at which cluster size threshold are calculated.\n(Press ENTER to accept default of .05 .02 .01 .005 .002 .001 .0005 .0002 .0001)\n"
      ii <- readline(prompt)
      if (ii == "") {
        athr <- ".05 .02 .01 .005 .002 .001 .0005 .0002 .0001"
      } else {
        # validate input
        athr <- check_nums(athr, lower = 1e-10, upper = .999)
      }
    }

    # TODO: It would be elegant to have gfeat_dir inputs to afni_3dclustsim, mirroring pTFCE... but it's not a priority at the moment.
    gfeat_list <- lapply(to_fwe$feat_dir, read_gfeat_dir)

    # every .gfeat folder may have multiple .feat subfolders (per cope)
    # thus, start with the doubly-nested list
    fwe_list <- lapply(gfeat_list, function(gg) {
      # iterate over cope.feat folders within the gfeat structure
      lapply(gg$cope_dirs, function(cc) {
        # convert internal default value to this mask
        if (clustsim_mask == "..DEFAULT..") {
          clustsim_mask <- cc$mask_file
        }

        afni_3dclustsim$new(
          residuals_file = cc$aux_files$res4d,
          residuals_mask_file = cc$mask_file,
          residuals_njobs = residual_njobs,
          ncpus = ncpus,
          prefix = "res4d",
          clustsim_mask = clustsim_mask,
          out_dir = dirname(cc$aux_files$res4d), # always put 3dClustSim outputs into same folder with res4d
          pthr = pthr, athr = athr
        )
      })
    })

    # unnest the .gfeat/.feat structure to just have a 1D list of clustsim objects
    fwe_list <- rlang::flatten(fwe_list)

    # build the clustsim objects into a list object that supports the $submit method (on all)
    clustsim_list <- afni_3dclustsim_list$new(obj_list = fwe_list)

    return(clustsim_list)
}

#' R6 class for an FWE correction method that can apply to one or more models
#'
#' @importFrom R6 R6Class
#' @importFrom data.table data.table
#' @importFrom checkmate assert_data_frame assert_logical
#' @keywords internal
fwe_spec <- R6::R6Class("fwe_spec",
  private = list(
    # data: keyed data.table object
    data = NULL,
    fwe_obj_list = list()
    #fwe_type = NULL
  ), public = list(
    #' @description create a new FWE correction instance
    #' @param ... not used yet
    #' @param fwe_data not used yet
    initialize = function(..., fwe_data = NULL) {
      objs <- list(...)
      #checkmate::assert_subset(fwe_type, c("ptfce, 3dclustsim", "palm", "randomise"), empty.ok = FALSE)

    },
    #' @description submit FWE estimation to the cluster
    submit = function() {
      # run submit method for every fwe object
      lapply(fwe_obj_list$submit())
    }
  )
)


#' Function to walk user through setting up FWE corrections for model outputs
#' @param gpa a \code{glm_pipeline_arguments} object containing model outputs in l3_model_setup
#' @param lg a logger object for logging messages
#' @return returns a modified gpa object that contains FWE specifications in $fwe_correction
#' @details currently only looks a level 3
#' @importFrom checkmate assert_class test_class
#' @export
build_fwe_correction <- function(gpa, lg = NULL) {
  checkmate::assert_class(gpa, "glm_pipeline_arguments")
  if (is.null(lg)) {
    lg <- lgr::get_logger("glm_pipeline/fwe_correction")
    lg$set_threshold(gpa$lgr_threshold)
  }
  checkmate::assert_class(lg, "Logger")

  level <- 3L # for now, only supporting FWE at group analysis level

  if (!checkmate::test_class(gpa$l3_models, "hi_model_set")) {
    msg <- "Cannot setup FWE correction until build_l3_models is complete and $l3_models is populated."
    lg$error(msg)
    stop(msg)
  }

  if (is.null(gpa$fwe_correction)) {
    gpa$fwe_correction <- list()
  }

  if ("fsl" %in% gpa$glm_software) {
    # enforce that at least some L3 models are complete
    enforce_glms_complete(gpa, level = 3L, lg)

    # complete_l3 <- #gpa$
  } else if ("spm" %in% gpa$glm_software) {
    stop("Not supported yet")
  } else if ("afni" %in% gpa$glm_software) {
    stop("Not supported yet")
  }

  # to setup an FWE specification, we only need the models in place, but they need not be complete.
  # cat(glue("\n\n---\nCurrent {field_desc} columns: {c_string(chosen_cols)}\n\n", .trim=FALSE))

  action <- 0L
  while (action != 4L) {
    action <- menu(c("Add FWE correction", "Modify FWE correction", "Delete FWE correction", "Done with FWE setup"),
      title = "What would you like to like to do?"
    )

    if (action == 1L) {
      fwe_name <- ""
      while (fwe_name == "") {
        prompt <- "Enter the unique name for this FWE correction setup: "
        nm <- readline(prompt)
        if (nm != "") {
          if (nm %in% names(gpa$fwe_correction)) {
            cat("FWE correction cannot have the same name as an existing correction!\n")
          } else {
            fwe_name <- nm
          }
        }
      }
      fwe_obj <- create_fwe_spec(gpa, level = level, lg = lg)
      gpa$fwe_correction[[fwe_name]] <- fwe_obj
    } else if (action == 2L) {
      cat("Not yet implemented!\n")
    } else if (action == 3L) {
      if (length(gpa$fwe_correction) == 0L) {
        message("Cannot delete an FWE correction because none exist!")
      } else {
        fwe_names <- names(gpa$fwe_correction)
        which_del <- menu(fwe_names, title = "Which FWE correction setup would you like to delete?")
        if (which_del > 0) {
          proceed <- menu(c("Proceed", "Cancel"), title = glue("Are you sure you want to delete {fwe_names[which_del]}?"))
          if (proceed == 1) {
            cat(glue("  Deleting {fwe_names[which_del]}\n"))
            gpa$fwe_correction[[which_del]] <- NULL
          } else {
            cat(glue("  Not deleting {fwe_names[which_del]}\n"))
          }
        }
      }
    } else if (action == 4L) {
      cat("Finishing FWE correction setup\n")
    }
  }

  return(gpa)
}

extract_glm_betas <- function(gpa, beta_definition = NULL) {
  checkmate::assert_class(gpa, "glm_pipeline_arguments")
  if (is.null(gpa$fwe_correction)) {
    warning("Cannot extract betas from GLMs using voxelwise thresholds because $fwe_correction is absent. Run build_fwe_correction() and run_fwe_correction() first!")
    return(invisible(NULL))
  }


}