make plot_min_depth_distribution work with ranger forests

Author: Yue-Jiang

DESCRIPTION

@@ -20,10 +20,11 @@ Imports:
   ggrepel (>= 0.6.5),
   MASS (>= 7.3.47),
   randomForest (>= 4.6.12),
+  ranger(>= 0.9.0),
   reshape2 (>= 1.4.2),
   rmarkdown (>= 1.5)
 Suggests:
   knitr
 VignetteBuilder: knitr
-RoxygenNote: 6.0.1
+RoxygenNote: 6.1.1
 URL: https://github.com/MI2DataLab/randomForestExplainer

NAMESPACE

@@ -1,5 +1,7 @@
 # Generated by roxygen2: do not edit by hand
 
+S3method(min_depth_distribution,randomForest)
+S3method(min_depth_distribution,ranger)
 export(explain_forest)
 export(important_variables)
 export(measure_importance)

R/min_depth_distribution.R

@@ -14,22 +14,43 @@ calculate_tree_depth <- function(frame){
   return(frame)
 }
 
+# Calculate the depth of each node in a single tree obtained from a forest with ranger::treeInfo
+calculate_tree_depth_ranger <- function(frame){
+  if(!all(c("rightChild", "leftChild") %in% names(frame))){
+    stop("The data frame has to contain columns called 'rightChild' and 'leftChild'!
+         It should be a product of the function ranger::treeInfo().")
+  }
+  frame$depth <- NA
+  frame$depth[1] <- 0
+  for(i in 2:nrow(frame)){
+    frame[i, "depth"] <-
+      frame[(!is.na(frame[, "leftChild"]) & frame[, "leftChild"] == frame[i, "nodeID"]) |
+              (!is.na(frame[, "rightChild"]) & frame[, "rightChild"] == frame[i, "nodeID"]), "depth"] + 1
+  }
+  return(frame)
+}
+
 #' Calculate minimal depth distribution of a random forest
 #'
 #' Get minimal depth values for all trees in a random forest
 #'
-#' @param forest A randomForest object
+#' @param forest A randomForest or ranger object
 #'
 #' @return A data frame with the value of minimal depth for every variable in every tree
 #'
-#' @import dplyr
-#' @importFrom data.table rbindlist
-#'
 #' @examples
 #' min_depth_distribution(randomForest::randomForest(Species ~ ., data = iris))
+#' min_depth_distribution(ranger::ranger(Species ~ ., data = iris))
 #'
 #' @export
 min_depth_distribution <- function(forest){
+  UseMethod("min_depth_distribution")
+}
+
+#' @import dplyr
+#' @importFrom data.table rbindlist
+#' @export
+min_depth_distribution.randomForest <- function(forest){
   tree <- NULL; `split var` <- NULL; depth <- NULL
   forest_table <-
     lapply(1:forest$ntree, function(i) randomForest::getTree(forest, k = i, labelVar = T) %>%
@@ -41,6 +62,21 @@ min_depth_distribution <- function(forest){
   return(min_depth_frame)
 }
 
+#' @import dplyr
+#' @importFrom data.table rbindlist
+#' @export
+min_depth_distribution.ranger <- function(forest){
+  tree <- NULL; splitvarName <- NULL; depth <- NULL
+  forest_table <-
+    lapply(1:forest$num.trees, function(i) ranger::treeInfo(forest, tree = i) %>%
+             calculate_tree_depth_ranger() %>% cbind(tree = i)) %>% rbindlist()
+  min_depth_frame <- dplyr::group_by(forest_table, tree, splitvarName) %>%
+    dplyr::summarize(min(depth))
+  colnames(min_depth_frame) <- c("tree", "variable", "minimal_depth")
+  min_depth_frame <- as.data.frame(min_depth_frame[!is.na(min_depth_frame$variable),])
+  return(min_depth_frame)
+}
+
 # Count the trees in which each variable had a given minimal depth
 min_depth_count <- function(min_depth_frame){
   tree <- NULL; minimal_depth <- NULL; variable <- NULL