Merge pull request #125 from tidymodels/faster-rf

Author: EmilHvitfeldt

NEWS.md

@@ -1,5 +1,7 @@
 # tidypredict (development version)
 
+- Speed up `tidypredict_fit()` for partykit and ranger packages. (#125)
+
 # tidypredict 0.5.1
 
 - Exported a number of internal functions to be used in {orbital} package

R/model-partykit.R

@@ -69,6 +69,9 @@ partykit_tree_info <- function(model) {
 get_pk_tree <- function(model) {
   tree <- partykit_tree_info(model)
   paths <- tree$nodeID[tree[, "terminal"]]
+
+  child_info <- get_child_info(tree)
+
   map(
     paths,
     ~ {
@@ -77,7 +80,7 @@ get_pk_tree <- function(model) {
       if (is.factor(prediction)) prediction <- as.character(prediction)
       list(
         prediction = prediction,
-        path = get_ra_path(.x, tree, FALSE)
+        path = get_ra_path(.x, tree, child_info, FALSE)
       )
     }
   )
@@ -132,6 +135,9 @@ tidypredict_fit.party <- function(model) {
 
   generate_one_tree <- function(tree_info) {
     paths <- tree_info$nodeID[tree_info[, "terminal"]]
+
+    child_info <- get_child_info(tree_info)
+
     paths <- map(
       paths,
       ~ {
@@ -140,7 +146,7 @@ tidypredict_fit.party <- function(model) {
           if (is.factor(prediction)) prediction <- as.character(prediction)
           list(
             prediction = prediction,
-            path = get_ra_path(.x, tree_info, FALSE)
+            path = get_ra_path(.x, tree_info, child_info, FALSE)
           )
        }
     )

R/model-ranger.R

@@ -1,33 +1,39 @@
 # Model parser ------------------------------------
-get_ra_path <- function(node_id, tree, default_op = TRUE) {
+get_ra_path <- function(node_id, tree, child_info, default_op = TRUE) {
   find <- node_id
   path <- node_id
-  for (j in node_id:0) {
-    row <- tree[tree$nodeID == j, ]
-    lc <- row["leftChild"][[1]] == find
-    lr <- row["rightChild"][[1]] == find
-    if (is.na(lc)) lc <- FALSE
-    if (is.na(lr)) lr <- FALSE
-    dir <- NULL
-    if (lc | lr) {
-      find <- j
-      path <- c(path, j)
+
+  leftChild <- tree$leftChild
+  rightChild <- tree$rightChild
+  splitval <- tree$splitval
+  splitclass <- tree$splitclass
+  splitvarName <- tree$splitvarName
+  
+  new <- child_info[[find]]
+  path <- find
+  repeat {
+    if (new == 0) {
+      path <- c(path, 0)
+      break
     }
+    path <- c(path, new)
+    find <- new
+    new <- child_info[[find]]
   }
+
   map2(
     path[1:length(path) - 1],
     path[2:length(path)],
     ~ {
-      rb <- tree[tree$nodeID == .y, ]
-      lc <- rb["leftChild"] == .x
-      lr <- rb["rightChild"] == .x
-      if (is.na(rb["splitval"][[1]])) {
+      lc <- leftChild[.y+1] == .x
+      lr <- rightChild[.y+1] == .x
+      if (is.na(splitval[.y+1])) {
         if (lc) op <- "in"
         if (lr) op <- "not-in"
-        vals <- strsplit(as.character(rb["splitclass"][[1]]), ", ")[[1]]
+        vals <- strsplit(as.character(splitclass[.y+1]), ", ")[[1]]
         list(
           type = "set",
-          col = as.character(rb["splitvarName"][[1]]),
+          col = as.character(splitvarName[.y+1]),
           vals = map(vals, ~.x),
           op = op
         )
@@ -41,18 +47,44 @@ get_ra_path <- function(node_id, tree, default_op = TRUE) {
         }
         list(
           type = "conditional",
-          col = as.character(rb["splitvarName"][[1]]),
-          val = rb["splitval"][[1]],
+          col = as.character(splitvarName[.y+1]),
+          val = splitval[.y+1],
           op = op
         )
       }
     }
   )
 }
 
+get_child_info <- function(tree) {
+  child_info <- numeric(max(tree$nodeID))
+  left_child <- tree$leftChild
+  right_child <- tree$rightChild
+  node_id <- tree$nodeID
+  
+  for (i in seq_len(nrow(tree))) {
+    node <- node_id[[i]]
+    
+    child <- left_child[[i]]
+    if (!is.na(child)) {
+      child_info[child] <- node
+    }
+    
+    child <- right_child[[i]]
+    if (!is.na(child)) {
+      child_info[child] <- node
+    }
+  }
+
+  child_info
+}
+
 get_ra_tree <- function(tree_no, model) {
   tree <- ranger::treeInfo(model, tree_no)
   paths <- tree$nodeID[tree[, "terminal"]]
+  
+  child_info <- get_child_info(tree)
+  
   map(
     paths,
     ~ {
@@ -61,7 +93,7 @@ get_ra_tree <- function(tree_no, model) {
         if (is.factor(prediction)) prediction <- as.character(prediction)
         list(
           prediction = prediction,
-          path = get_ra_path(.x, tree, TRUE)
+          path = get_ra_path(.x, tree, child_info, TRUE)
         )
       } else {
         preds <- map_lgl(colnames(tree), ~ "pred." == substr(.x, 1, 5))
@@ -79,7 +111,7 @@ get_ra_tree <- function(tree_no, model) {
           prediction = prediction,
           prob = prob,
           probs = predictions,
-          path = get_ra_path(.x, tree, TRUE)
+          path = get_ra_path(.x, tree, child_info, TRUE)
         )
       }
     }

tests/testthat/_snaps/model-cubist.md

@@ -17,7 +17,7 @@ run_test <- function(model, test_formula = TRUE) {
   }
 }
 
-run_test(
+tidypredict_fit(
   ranger::ranger(Species ~ ., data = iris, num.trees = num_trees, seed = 100, num.threads = 2)
 )
 

tests/testthat/test-model-ranger.R

@@ -1,4 +1,5 @@
 Version: 1.0
+ProjectId: 0cda5e71-1965-4603-b39a-db0e57a2cde6
 
 RestoreWorkspace: Default
 SaveWorkspace: Default