Initial working version of leaf index sparse matrix representation

Author: andrewherren

include/stochtree/ensemble.h

@@ -220,6 +220,45 @@ class TreeEnsemble {
     }
   }
 
+  /*!
+   * \brief Obtain a 0-based leaf index for every tree in an ensemble and for each 
+   *        observation in a ForestDataset. Internally, trees are stored as essentially 
+   *        vectors of node information, and the leaves_ vector gives us node IDs for every 
+   *        leaf in the tree. Here, we would like to know, for every observation in a dataset, 
+   *        which leaf number it is mapped to. Since the leaf numbers themselves 
+   *        do not carry any information, we renumber them from 0 to `leaves_.size()-1`. 
+   *        We compute this at the tree-level and coordinate this computation at the 
+   *        ensemble level.
+   *
+   *        Note: this assumes the creation of a vector of column indices of size 
+   *        `dataset.NumObservations()` x `ensemble.NumTrees()`
+   * \param ForestDataset Dataset with which to predict leaf indices from the tree
+   * \param output Vector of length num_trees*n which stores the leaf node prediction
+   * \param num_trees Number of trees in an ensemble
+   * \param n Size of dataset
+   */
+  void PredictLeafIndicesInplace(ForestDataset* dataset, std::vector<int32_t>& output, int num_trees, data_size_t n) {
+    CHECK_GE(output.size(), num_trees*n);
+    int offset = 0;
+    for (int j = 0; j < num_trees; j++) {
+      auto &tree = *trees_[j];
+      tree.PredictLeafIndexInplace(dataset, output, offset);
+      offset += n;
+    }
+  }
+
+  /*!
+   * \brief Same as `PredictLeafIndicesInplace` but assumes responsibility for allocating and returning output vector.
+   * \param ForestDataset Dataset with which to predict leaf indices from the tree
+   */
+  std::vector<int32_t> PredictLeafIndices(ForestDataset* dataset) {
+    int num_trees = num_trees_;
+    data_size_t n = dataset->NumObservations();
+    std::vector<int32_t> output(n*num_trees);
+    PredictLeafIndicesInplace(dataset, output, num_trees, n);
+    return output;
+  }
+
   /*! \brief Save to JSON */
   json to_json() {
     json result_obj;

include/stochtree/tree.h

@@ -7,6 +7,7 @@
 #define STOCHTREE_TREE_H_
 
 #include <nlohmann/json.hpp>
+#include <stochtree/data.h>
 #include <stochtree/log.h>
 #include <stochtree/meta.h>
 #include <Eigen/Dense>
@@ -586,6 +587,25 @@ class Tree {
    */
   void SetLeafVector(std::int32_t nid, std::vector<double> const& leaf_vector);
 
+  /*!
+   * \brief Obtain a 0-based leaf index for each observation in a ForestDataset.
+   *        Internally, trees are stored as essentially vectors of node information, 
+   *        and the leaves_ vector gives us node IDs for every leaf in the tree.
+   *        Here, we would like to know, for every observation in a dataset, 
+   *        which leaf number it is mapped to. Since the leaf numbers themselves 
+   *        do not carry any information, we renumber them from 0 to `leaves_.size()-1`. 
+   *
+   *        Note: this is a tree-level helper function for an ensemble-level function. 
+   *        It assumes the creation of: 
+   *           (a) a vector of column indices of size `dataset.NumObservations()` x `ensemble.NumTrees()`, and
+   *           (b) a running counter of the number of tree-observations already indexed in the ensemble  
+   *               (used as offsets for the leaf number computed and returned here)
+   *        Users running this function for a single tree may simply pre-allocate an output vector as 
+   *        std::vector<int32_t> output(dataset->NumObservations()) and set the offset to 0.
+   * \param dataset Dataset with which to predict leaf indices from the tree
+   */
+  void PredictLeafIndexInplace(ForestDataset* dataset, std::vector<int32_t>& output, int32_t offset);
+
   // Node info
   std::vector<TreeNodeType> node_type_;
   std::vector<std::int32_t> parent_;

src/tree.cpp

@@ -420,6 +420,20 @@ void Tree::SetLeafVector(std::int32_t nid, std::vector<double> const& node_leaf_
   node_type_.at(nid) = TreeNodeType::kLeafNode;
 }
 
+void Tree::PredictLeafIndexInplace(ForestDataset* dataset, std::vector<int32_t>& output, int32_t offset) {
+  int n = dataset->NumObservations();
+  CHECK_GE(output.size(), offset + n);
+  std::map<int32_t,int32_t> renumber_map;
+  for (int i = 0; i < leaves_.size(); i++) {
+    renumber_map.insert({leaves_[i], i});
+  }
+  int32_t node_id;
+  for (int i = 0; i < n; i++) {
+    node_id = EvaluateTree(*this, dataset->GetCovariates(), i);
+    output.at(offset + i) = renumber_map.at(node_id);
+  }
+}
+
 void TreeNodeVectorsToJson(json& obj, Tree* tree) {
   // Initialize a map with names of the node vectors and empty json arrays
   std::map<std::string, json> tree_array_map;

test/test_tree.cpp

@@ -4,6 +4,7 @@
  */
 #include <gtest/gtest.h>
 #include <testutils.h>
+#include <stochtree/data.h>
 #include <stochtree/log.h>
 #include <stochtree/tree.h>
 #include <iostream>
@@ -153,3 +154,25 @@ TEST(Tree, MultivariateTreeCategoricalSplitConstruction) {
   ASSERT_TRUE(tree.IsLeaf(1));
   ASSERT_TRUE(tree.IsLeaf(2));
 }
+
+TEST(Tree, SparseLeafRepresentation) {
+  // Construct small tree
+  StochTree::Tree tree;
+  tree.Init(1);
+  tree.ExpandNode(0, 0, 0.5, 0., 0.);
+  
+  // Load test data
+  StochTree::TestUtils::TestDataset test_dataset;
+  test_dataset = StochTree::TestUtils::LoadSmallDatasetUnivariateBasis();
+
+  // Construct datasets
+  int n = test_dataset.n;
+  StochTree::ForestDataset dataset = StochTree::ForestDataset();
+  dataset.AddCovariates(test_dataset.covariates.data(), n, test_dataset.x_cols, test_dataset.row_major);
+
+  // Predict leaf indices of each observation in `dataset`
+  std::vector<int32_t> leaf_index_preds(n);
+  tree.PredictLeafIndexInplace(&dataset, leaf_index_preds, 0);
+  std::vector<int32_t> leaf_index_expected{1,1,0,1,1,1,1,1,0,0};
+  ASSERT_EQ(leaf_index_expected, leaf_index_preds);
+}