Merge pull request #34 from StochasticTree/prediction_cache

Added tree prediction caching to the ForestTracker

Author: andrewherren

include/stochtree/partition_tracker.h

@@ -41,6 +41,7 @@ namespace StochTree {
 
 /*! \brief Forward declarations of component classes */
 class SampleNodeMapper;
+class SamplePredMapper;
 class UnsortedNodeSampleTracker;
 class SortedNodeSampleTracker;
 class FeaturePresortRootContainer;
@@ -52,9 +53,13 @@ class ForestTracker {
   ~ForestTracker() {}
   void AssignAllSamplesToRoot();
   void AssignAllSamplesToRoot(int32_t tree_num);
+  void AssignAllSamplesToConstantPrediction(double value);
+  void AssignAllSamplesToConstantPrediction(int32_t tree_num, double value);
   void ResetRoot(Eigen::MatrixXd& covariates, std::vector<FeatureType>& feature_types, int32_t tree_num);
   void AddSplit(Eigen::MatrixXd& covariates, TreeSplit& split, int32_t split_feature, int32_t tree_id, int32_t split_node_id, int32_t left_node_id, int32_t right_node_id, bool keep_sorted = false);
   void RemoveSplit(Eigen::MatrixXd& covariates, Tree* tree, int32_t tree_id, int32_t split_node_id, int32_t left_node_id, int32_t right_node_id, bool keep_sorted = false);
+  double GetTreeSamplePrediction(data_size_t sample_id, int tree_id);
+  void SetTreeSamplePrediction(data_size_t sample_id, int tree_id, double value);
   data_size_t GetNodeId(int observation_num, int tree_num);
   data_size_t UnsortedNodeBegin(int tree_id, int node_id);
   data_size_t UnsortedNodeEnd(int tree_id, int node_id);
@@ -66,11 +71,14 @@ class ForestTracker {
   std::vector<data_size_t>::iterator UnsortedNodeEndIterator(int tree_id, int node_id);
   std::vector<data_size_t>::iterator SortedNodeBeginIterator(int node_id, int feature_id);
   std::vector<data_size_t>::iterator SortedNodeEndIterator(int node_id, int feature_id);
+  SamplePredMapper* GetSamplePredMapper() {return sample_pred_mapper_.get();}
   SampleNodeMapper* GetSampleNodeMapper() {return sample_node_mapper_.get();}
   UnsortedNodeSampleTracker* GetUnsortedNodeSampleTracker() {return unsorted_node_sample_tracker_.get();}
   SortedNodeSampleTracker* GetSortedNodeSampleTracker() {return sorted_node_sample_tracker_.get();}
 
  private:
+  /*! \brief Mapper from observations to predicted values for every tree in a forest */
+  std::unique_ptr<SamplePredMapper> sample_pred_mapper_;
   /*! \brief Mapper from observations to leaf node indices for every tree in a forest */
   std::unique_ptr<SampleNodeMapper> sample_node_mapper_;
   /*! \brief Data structure tracking / updating observations available in each node for every tree in a forest
@@ -88,6 +96,47 @@ class ForestTracker {
   int num_features_;
 };
 
+/*! \brief Class storing sample-prediction map for each tree in an ensemble */
+class SamplePredMapper {
+ public:
+  SamplePredMapper(int num_trees, data_size_t num_observations) {
+    num_trees_ = num_trees;
+    num_observations_ = num_observations;
+    // Initialize the vector of vectors of leaf indices for each tree
+    tree_preds_.resize(num_trees_);
+    for (int j = 0; j < num_trees_; j++) {
+      tree_preds_[j].resize(num_observations_);
+    }
+  }
+
+  inline double GetPred(data_size_t sample_id, int tree_id) {
+    CHECK_LT(sample_id, num_observations_);
+    CHECK_LT(tree_id, num_trees_);
+    return tree_preds_[tree_id][sample_id];
+  }
+
+  inline void SetPred(data_size_t sample_id, int tree_id, double value) {
+    CHECK_LT(sample_id, num_observations_);
+    CHECK_LT(tree_id, num_trees_);
+    tree_preds_[tree_id][sample_id] = value;
+  }
+  
+  inline int NumTrees() {return num_trees_;}
+  
+  inline int NumObservations() {return num_observations_;}
+
+  inline void AssignAllSamplesToConstantPrediction(int tree_id, double value) {
+    for (data_size_t i = 0; i < num_observations_; i++) {
+      tree_preds_[tree_id][i] = value;
+    }
+  }
+
+ private:
+  std::vector<std::vector<double>> tree_preds_;
+  int num_trees_;
+  data_size_t num_observations_;
+};
+
 /*! \brief Class storing sample-node map for each tree in an ensemble */
 class SampleNodeMapper {
  public:

include/stochtree/tree_sampler.h

@@ -162,18 +162,28 @@ static double ComputeMeanOutcome(ColumnVector& residual) {
   return total_outcome / static_cast<double>(n);
 }
 
-static void UpdateResidualTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, bool requires_basis, std::function<double(double, double)> op) {
+static void UpdateResidualTree(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, Tree* tree, int tree_num, bool requires_basis, std::function<double(double, double)> op, bool tree_new) {
   data_size_t n = dataset.GetCovariates().rows();
   double pred_value;
   int32_t leaf_pred;
   double new_resid;
   for (data_size_t i = 0; i < n; i++) {
-    leaf_pred = tracker.GetNodeId(i, tree_num);
-    if (requires_basis) {
-      pred_value = tree->PredictFromNode(leaf_pred, dataset.GetBasis(), i);
+    if (tree_new) {
+      // If the tree has been newly sampled or adjusted, we must rerun the prediction 
+      // method and update the SamplePredMapper stored in tracker
+      leaf_pred = tracker.GetNodeId(i, tree_num);
+      if (requires_basis) {
+        pred_value = tree->PredictFromNode(leaf_pred, dataset.GetBasis(), i);
+      } else {
+        pred_value = tree->PredictFromNode(leaf_pred);
+      }
+      tracker.SetTreeSamplePrediction(i, tree_num, pred_value);
     } else {
-      pred_value = tree->PredictFromNode(leaf_pred);
+      // If the tree has not yet been modified via a sampling step, 
+      // we can query its prediction directly from the SamplePredMapper stored in tracker
+      pred_value = tracker.GetTreeSamplePrediction(i, tree_num);
     }
+    // Run op (either plus or minus) on the residual and the new prediction
     new_resid = op(residual.GetElement(i), pred_value);
     residual.SetElement(i, new_resid);
   }
@@ -210,7 +220,7 @@ class MCMCForestSampler {
     for (int i = 0; i < num_trees; i++) {
       // Add tree i's predictions back to the residual (thus, training a model on the "partial residual")
       tree = ensemble->GetTree(i);
-      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), plus_op_);
+      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), plus_op_, false);
 
       // Sample tree i
       tree = ensemble->GetTree(i);
@@ -222,7 +232,7 @@ class MCMCForestSampler {
 
       // Subtract tree i's predictions back out of the residual
       tree = ensemble->GetTree(i);
-      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), minus_op_);
+      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), minus_op_, true);
     }
   }
 
@@ -477,7 +487,7 @@ class GFRForestSampler {
     for (int i = 0; i < num_trees; i++) {
       // Add tree i's predictions back to the residual (thus, training a model on the "partial residual")
       Tree* tree = ensemble->GetTree(i);
-      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), plus_op_);
+      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), plus_op_, false);
 
       // Reset the tree and sample trackers
       ensemble->ResetInitTree(i);
@@ -492,7 +502,7 @@ class GFRForestSampler {
       leaf_model.SampleLeafParameters(dataset, tracker, residual, tree, i, global_variance, gen);
 
       // Subtract tree i's predictions back out of the residual
-      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), minus_op_);
+      UpdateResidualTree(tracker, dataset, residual, tree, i, leaf_model.RequiresBasis(), minus_op_, true);
     }
   }
 

src/partition_tracker.cpp

@@ -15,6 +15,7 @@
 namespace StochTree {
 
 ForestTracker::ForestTracker(Eigen::MatrixXd& covariates, std::vector<FeatureType>& feature_types, int num_trees, int num_observations) {
+  sample_pred_mapper_ = std::make_unique<SamplePredMapper>(num_trees, num_observations);
   sample_node_mapper_ = std::make_unique<SampleNodeMapper>(num_trees, num_observations);
   unsorted_node_sample_tracker_ = std::make_unique<UnsortedNodeSampleTracker>(num_observations, num_trees);
   presort_container_ = std::make_unique<FeaturePresortRootContainer>(covariates, feature_types);
@@ -69,6 +70,16 @@ void ForestTracker::AssignAllSamplesToRoot(int32_t tree_num) {
   sample_node_mapper_->AssignAllSamplesToRoot(tree_num);
 }
 
+void ForestTracker::AssignAllSamplesToConstantPrediction(double value) {
+  for (int i = 0; i < num_trees_; i++) {
+    sample_pred_mapper_->AssignAllSamplesToConstantPrediction(i, value);
+  }
+}
+
+void ForestTracker::AssignAllSamplesToConstantPrediction(int32_t tree_num, double value) {
+  sample_pred_mapper_->AssignAllSamplesToConstantPrediction(tree_num, value);
+}
+
 void ForestTracker::AddSplit(Eigen::MatrixXd& covariates, TreeSplit& split, int32_t split_feature, int32_t tree_id, int32_t split_node_id, int32_t left_node_id, int32_t right_node_id, bool keep_sorted) {
   sample_node_mapper_->AddSplit(covariates, split, split_feature, tree_id, split_node_id, left_node_id, right_node_id);
   unsorted_node_sample_tracker_->PartitionTreeNode(covariates, tree_id, split_node_id, left_node_id, right_node_id, split_feature, split);
@@ -83,6 +94,14 @@ void ForestTracker::RemoveSplit(Eigen::MatrixXd& covariates, Tree* tree, int32_t
   // TODO: WARN if this is called from the GFR Tree Sampler
 }
 
+double ForestTracker::GetTreeSamplePrediction(data_size_t sample_id, int tree_id) {
+  return sample_pred_mapper_->GetPred(sample_id, tree_id);
+}
+
+void ForestTracker::SetTreeSamplePrediction(data_size_t sample_id, int tree_id, double value) {
+  sample_pred_mapper_->SetPred(sample_id, tree_id, value);
+}
+
 FeatureUnsortedPartition::FeatureUnsortedPartition(data_size_t n) {
   indices_.resize(n);
   std::iota(indices_.begin(), indices_.end(), 0);