R/plot.gg_vimp.R

In ggRandomForests: Visually Exploring Random Forests

####**********************************************************************
####**********************************************************************
####
####  ----------------------------------------------------------------
####  Written by:
####  ----------------------------------------------------------------
####    John Ehrlinger, Ph.D.
####
####    email:  john.ehrlinger@gmail.com
####    URL:    https://github.com/ehrlinger/ggRandomForests
####  ----------------------------------------------------------------
####
####**********************************************************************
####**********************************************************************
#' Plot a \code{\link{gg_vimp}} object, extracted variable importance of a
#' \code{\link[randomForestSRC]{rfsrc}} object
#'
#' Draws a horizontal bar chart of the VIMP scores extracted by
#' \code{\link{gg_vimp}}.  Each bar represents one predictor; bar length is
#' proportional to its permutation VIMP -- the average rise in OOB prediction
#' error when that predictor's OOB values are randomly shuffled.  Predictors
#' are sorted in descending order of importance so the most influential
#' variables appear at the top.
#'
#' Bars are colored by the \code{positive} flag: a bar at or below zero
#' (non-positive VIMP) is color-coded differently to flag predictors that
#' \emph{hurt} OOB accuracy when their signal is removed -- usually a sign of
#' collinearity or a very noisy variable.  In a well-behaved forest most bars
#' are positive; the color distinction matters when a handful are not.
#'
#' @param x \code{\link{gg_vimp}} object created from a
#' \code{\link[randomForestSRC]{rfsrc}} object
#' @param relative should we plot vimp or relative vimp. Defaults to vimp.
#' @param lbls A vector of alternative variable labels. Item names should be
#' the same as the variable names.
#' @param ... optional arguments passed to gg_vimp if necessary
#'
#' @return \code{ggplot} object
#'
#' @seealso \code{\link{gg_vimp}}
#'
#' @references
#' Breiman L. (2001). Random forests, Machine Learning, 45:5-32.
#'
#' Ishwaran H. and Kogalur U.B. (2007). Random survival forests for
#' R, Rnews, 7(2):25-31.
#'
#' Ishwaran H. and Kogalur U.B. randomForestSRC: Random Forests for Survival,
#' Regression and Classification. R package version >= 3.4.0.
#' \url{https://cran.r-project.org/package=randomForestSRC}
#'
#' @examples
#' ## ------------------------------------------------------------
#' ## classification example
#' ## ------------------------------------------------------------
#' ## -------- iris data
#' rfsrc_iris <- randomForestSRC::rfsrc(Species ~ ., data = iris)
#' gg_dta <- gg_vimp(rfsrc_iris)
#' plot(gg_dta)
#'
#' ## ------------------------------------------------------------
#' ## regression example
#' ## ------------------------------------------------------------
#' ## -------- air quality data
#' rfsrc_airq <- randomForestSRC::rfsrc(Ozone ~ ., airquality)
#' gg_dta <- gg_vimp(rfsrc_airq)
#' plot(gg_dta)
#'
#'
#' @export
plot.gg_vimp <- function(x, relative, lbls, ...) {
  gg_dta <- x

  # Accept raw rfsrc / randomForest objects and compute VIMP on the fly
  if (!inherits(gg_dta, "gg_vimp")) {
    gg_dta <- gg_vimp(gg_dta, ...)
  }

  # Capture extra args so we can inspect nvar.
  arg_set <- list(...)

  # Optionally restrict to the top-nvar most important variables (gg_vimp
  # already sorts by descending VIMP, so we just trim the tail).
  nvar <- nrow(gg_dta)
  if (!is.null(arg_set$nvar)) {
    if (is.numeric(arg_set$nvar) && arg_set$nvar > 1) {
      if (arg_set$nvar < nrow(gg_dta)) {
        nvar <- arg_set$nvar
        gg_dta <- gg_dta[seq_len(nvar), ]
      }
    }
  }

  gg_plt <- ggplot2::ggplot(gg_dta)

  # Use "vimp" as the bar-height column when it exists; fall back to the
  # first column name for objects that store a renamed importance measure.
  msr <- "vimp"
  if (!msr %in% colnames(gg_dta)) {
    msr <- colnames(gg_dta)[1]
  }

  # Always map both `fill` and `color` to `positive` -- this gives filled bars
  # (rather than hollow outlines) and ensures the function-level
  # `labs(fill = ..., color = ...)` below applies cleanly. When `positive` has
  # only one level (all VIMPs positive, the common case for well-behaved
  # forests), the bars simply render in a single color and ggplot collapses
  # the fill+color legend into a one-row legend; when both signs are present,
  # the two-row legend distinguishes positive from negative VIMP.
  #
  # Both aesthetics share the same legend title ("VIMP > 0") so ggplot
  # collapses what would otherwise be two legends -- one for fill and one
  # titled with the column name "positive" -- into a single merged legend.
  legend_title <- "VIMP > 0"
  gg_plt <- gg_plt +
    ggplot2::geom_bar(
      ggplot2::aes(
        y    = .data[[msr]],
        x    = .data$vars,
        fill = .data$positive,
        color = .data$positive
      ),
      stat = "identity",
      width = .5,
    )
  # Set both legends' titles to the same string so ggplot merges them.
  # Users can override with their own labs() call after the fact.
  gg_plt <- gg_plt +
    ggplot2::labs(x = "", y = msr,
                  fill = legend_title, color = legend_title)

  if (!missing(lbls)) {
    # Map internal variable names to human-readable labels.  lbls should be a
    # named character vector; any unmatched variables keep their original name.
    if (length(lbls) >= length(gg_dta$vars)) {
      st_lbls <- lbls[as.character(gg_dta$vars)]
      names(st_lbls) <- as.character(gg_dta$vars)
      # Fall back to the raw variable name when no label was supplied
      st_lbls[which(is.na(st_lbls))] <-
        names(st_lbls[which(is.na(st_lbls))])

      gg_plt <- gg_plt +
        ggplot2::scale_x_discrete(labels = st_lbls)
    }
  }

  # Flip coordinates so variable names appear on the y-axis (horizontal bars
  # are easier to read when there are many variables).  If gg_vimp contains
  # a "set" column (comparison VIMP across two forests), facet by set.
  if (is.null(gg_dta$set) || length(unique(gg_dta$set)) < 2) {
    gg_plt <- gg_plt +
      ggplot2::coord_flip()
  } else {
    gg_plt <- gg_plt +
      ggplot2::coord_flip() +
      ggplot2::facet_grid(~set)
  }

  return(gg_plt)
}