Predict same leaf labels as printed labels with hybrid policytree (#153)

Author: erikcs

r-package/policytree/R/hybrid_policy_tree.R

@@ -183,7 +183,7 @@ unpack_tree <- function(tree) {
 # For "largish" depth it could also be sparse.
 tree_mat <- function(nodes, depth) {
   num.nodes <- 2^(depth + 1) - 1
-  tree.array <- matrix(0, num.nodes, 4)
+  tree.array <- matrix(0, num.nodes, 5)
   frontier <- 1
   i <- 1
   j <- 1
@@ -193,11 +193,13 @@ tree_mat <- function(nodes, depth) {
     if (nodes[[node]]$is_leaf) {
       tree.array[j, 1] <- -1
       tree.array[j, 2] <- nodes[[node]]$action
+      tree.array[j, 5] <- node - 1
     } else {
       tree.array[j, 1] <- nodes[[node]]$split_variable
       tree.array[j, 2] <- nodes[[node]]$split_value
       tree.array[j, 3] <- i + 1
       tree.array[j, 4] <- i + 2
+      tree.array[j, 5] <- node - 1
       frontier <- c(frontier, nodes[[node]]$left_child, nodes[[node]]$right_child)
       i <- i + 2
     }

r-package/policytree/src/Rcppbindings.cpp

@@ -58,10 +58,10 @@ Rcpp::List tree_search_rcpp(const Rcpp::NumericMatrix& X,
   // We store the tree as the same list data structure (`nodes`) as GRF for seamless integration with
   // the plot and print methods. We also store the tree as an array (`tree_array`) for faster lookups.
   // This will make a difference for a very large amount of lookups, like n = 1 000 000.
-  // The columns 0 to 3 are:
-  // split_variable (-1 if leaf) | split_value (action_id if leaf) | left_child | right_child
+  // The columns 0 to 4 are:
+  // split_variable (-1 if leaf) | split_value (action_id if leaf) | left_child | right_child | optional node_id
   int num_nodes = pow(2.0, depth + 1.0) - 1;
-  Rcpp::NumericMatrix tree_array(num_nodes, 4);
+  Rcpp::NumericMatrix tree_array(num_nodes, 5);
   Rcpp::List nodes;
   int i = 1;
   int j = 0;
@@ -76,6 +76,7 @@ Rcpp::List tree_search_rcpp(const Rcpp::NumericMatrix& X,
       nodes.push_back(list_node);
       tree_array(j, 0) = -1;
       tree_array(j, 1) = node->action_id + 1;
+      tree_array(j, 4) = j; // only used by hybrid policytree to predict leaf number that matches up with the printed leaf number.
     } else {
       auto list_node = Rcpp::List::create(Rcpp::Named("is_leaf") = false,
                                           Rcpp::Named("split_variable") = node->index + 1, // C++ index
@@ -87,6 +88,7 @@ Rcpp::List tree_search_rcpp(const Rcpp::NumericMatrix& X,
       tree_array(j, 1) = node->value;
       tree_array(j, 2) = i + 1; // left child
       tree_array(j, 3) = i + 2; // right child
+      tree_array(j, 4) = j;
       frontier.push(std::move(node->left_child));
       frontier.push(std::move(node->right_child));
       i += 2;
@@ -120,8 +122,9 @@ Rcpp::NumericMatrix tree_search_rcpp_predict(const Rcpp::NumericMatrix& tree_arr
       bool is_leaf = tree_array(node, 0) == -1;
       if (is_leaf) {
         size_t action = tree_array(node, 1);
+        size_t node_id = tree_array(node, 4);
         result(sample, 0) = action;
-        result(sample, 1) = node;
+        result(sample, 1) = node_id;
         break;
       }
       size_t split_var = tree_array(node, 0) - 1; // Offset by 1 for C++ indexing

r-package/policytree/tests/testthat/test_hybrid_policy_tree.R

@@ -1,7 +1,7 @@
 test_that("hybrid_policy_tree works as expected", {
   n <- 500
   p <- 2
-  d <- 2
+  d <- 42
   X <- round(matrix(rnorm(n * p), n, p), 1)
   Y <- matrix(runif(n * d), n, d)
 
@@ -25,6 +25,17 @@ test_that("hybrid_policy_tree works as expected", {
     }
   }
   expect_equal(hpp, pp)
+  # is internally consistent with predicted node ids.
+  hpp.node <- predict(htree, X, type = "node.id")
+  values <- aggregate(Y, by = list(leaf.node = hpp.node), FUN = mean)
+  best <- apply(values[, -1], 1, FUN = which.max)
+  expect_equal(best[match(hpp.node, values[, 1])], hpp)
+
+  # uses node labels that are the same as the printed labels.
+  printed.node.id <- lapply(seq_along(htree$nodes), function(i) {
+    if (htree$nodes[[i]]$is_leaf) i
+  })
+  expect_true(all(unique(hpp.node) %in% printed.node.id))
 
   # search.depth = 1 when a single split is optimal is identical to a depth 1 policy_tree
   n <- 250