Fix ranger.unify for probability forests (#43)

* unify_ranger.R: fix for probability forests

* Update unify_ranger.R

* Update unify_ranger.R

Author: ck37

R/unify_ranger.R

@@ -40,6 +40,10 @@ ranger.unify <- function(rf_model, data) {
   n <- rf_model$num.trees
   x <- lapply(1:n, function(tree) {
     tree_data <- data.table::as.data.table(ranger::treeInfo(rf_model, tree = tree))
+    # Fix for probability forests
+    if (rf_model$treetype == "Probability estimation") {
+      data.table::setnames(tree_data, "pred.1", "prediction")
+    }
     tree_data[, c("nodeID",  "leftChild", "rightChild", "splitvarName", "splitval", "prediction")]
   })
   return(ranger_unify.common(x = x, n = n, data = data, feature_names = rf_model$forest$independent.variable.names))