R/variable_importance.R

In juoe/sdmflow: Flexible Workflows for Species Distribution Modeling

#' Get variable importance from data frame with models objects.
#'
#' @param data A data frame with column "model" containing the model objects.
#' @param long Boolean. Whether output should be as "long" table suited for plotting. Default is FALSE.
#' @param keep_vars Optional character of variables to keep while unnesting.
#'
#' @return Input data frame with added column "var_imp" containing relative variable importance.
#' @export


variable_importance <- function(data, long = FALSE, keep_vars=NULL, ...) {

  out <- dplyr::mutate(data, var_imp = purrr::map(model , function(model) {
    if(class(model) == "gbm"){
      gbm::relative.influence(model, model$n.trees, ...) %>%
        as.data.frame() %>%
        setNames("imp") %>%
        mutate(var = rownames(.)) %>%
        setNames(c("imp","var"))
    } else {
      caret::varImp(model, ...) %>% mutate(var = rownames(.)) %>% setNames(c("imp","var"))
    }

  }))

  if(long){
    out <- suppressWarnings(tidyr::unnest(out, var_imp) %>% dplyr::select(dplyr::one_of(c("imp",
                                                                                           "var",
                                                                                           "data_name",
                                                                                           "predictor_name",
                                                                                           "sampling_name",
                                                                                           "resampling_name",
                                                                                           "model_name",
                                                                                          keep_vars))))
  }

  return(out)
}