Merge pull request #35 from StochasticTree/expand_initialization_api

Expand initialization api

Author: andrewherren

include/stochtree/container.h

@@ -26,6 +26,8 @@ class ForestContainer {
   ForestContainer(int num_samples, int num_trees, int output_dimension = 1, bool is_leaf_constant = true);
   ~ForestContainer() {}
 
+  void InitializeRoot(double leaf_value);
+  void InitializeRoot(std::vector<double>& leaf_vector);
   void AddSamples(int num_samples);
   void CopyFromPreviousSample(int new_sample_id, int previous_sample_id);
   std::vector<double> Predict(ForestDataset& dataset);
@@ -39,8 +41,13 @@ class ForestContainer {
   inline int32_t NumLeaves(int ensemble_num) {return forests_[ensemble_num]->NumLeaves();}
   inline int32_t OutputDimension() {return output_dimension_;}
   inline int32_t OutputDimension(int ensemble_num) {return forests_[ensemble_num]->OutputDimension();}
+  inline bool IsLeafConstant() {return is_leaf_constant_;}
   inline bool IsLeafConstant(int ensemble_num) {return forests_[ensemble_num]->IsLeafConstant();}
-  
+  inline bool AllRoots(int ensemble_num) {return forests_[ensemble_num]->AllRoots();}
+  inline void SetLeafValue(int ensemble_num, double leaf_value) {forests_[ensemble_num]->SetLeafValue(leaf_value);}
+  inline void SetLeafVector(int ensemble_num, std::vector<double>& leaf_vector) {forests_[ensemble_num]->SetLeafVector(leaf_vector);}
+  inline void IncrementSampleCount() {num_samples_++;}
+
   void SaveToJsonFile(std::string filename) {
     nlohmann::json model_json = this->to_json();
     std::ofstream output_file(filename);

include/stochtree/ensemble.h

@@ -195,6 +195,31 @@ class TreeEnsemble {
     return is_leaf_constant_;
   }
 
+  inline bool AllRoots() {
+    for (int i = 0; i < num_trees_; i++) {
+      if (!trees_[i]->IsRoot()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  inline void SetLeafValue(double leaf_value) {
+    CHECK_EQ(output_dimension_, 1);
+    for (int i = 0; i < num_trees_; i++) {
+      CHECK(trees_[i]->IsRoot());
+      trees_[i]->SetLeaf(0, leaf_value);
+    }
+  }
+
+  inline void SetLeafVector(std::vector<double>& leaf_vector) {
+    CHECK_EQ(output_dimension_, leaf_vector.size());
+    for (int i = 0; i < num_trees_; i++) {
+      CHECK(trees_[i]->IsRoot());
+      trees_[i]->SetLeafVector(0, leaf_vector);
+    }
+  }
+
   /*! \brief Save to JSON */
   json to_json() {
     json result_obj;

include/stochtree/tree.h

@@ -95,6 +95,9 @@ class Tree {
   void ExpandNode(std::int32_t nid, int split_index, TreeSplit& split, double left_value, double right_value);
   /*! \brief Expand a node based on a generic split rule */
   void ExpandNode(std::int32_t nid, int split_index, TreeSplit& split, std::vector<double> left_value_vector, std::vector<double> right_value_vector);
+
+  /*! \brief Whether or not a tree is a "stump" consisting of a single root node */
+  inline bool IsRoot() {return leaves_.size() == 1;}
 
   /*! \brief Save to JSON */
   json to_json();

include/stochtree/tree_sampler.h

@@ -162,6 +162,31 @@ static double ComputeMeanOutcome(ColumnVector& residual) {
   return total_outcome / static_cast<double>(n);
 }
 
+static void UpdateResidualEntireForest(ForestTracker& tracker, ForestDataset& dataset, ColumnVector& residual, TreeEnsemble* forest, bool requires_basis, std::function<double(double, double)> op) {
+  data_size_t n = dataset.GetCovariates().rows();
+  double tree_pred = 0.;
+  double pred_value = 0.;
+  double new_resid = 0.;
+  int32_t leaf_pred;
+  for (data_size_t i = 0; i < n; i++) {
+    for (int j = 0; j < forest->NumTrees(); j++) {
+      Tree* tree = forest->GetTree(j);
+      leaf_pred = tracker.GetNodeId(i, j);
+      if (requires_basis) {
+        tree_pred += tree->PredictFromNode(leaf_pred, dataset.GetBasis(), i);
+      } else {
+        tree_pred += tree->PredictFromNode(leaf_pred);
+      }
+      tracker.SetTreeSamplePrediction(i, j, tree_pred);
+      pred_value += tree_pred;
+    }
+    
+    // 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);
+  }
+}
+
 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;
@@ -196,21 +221,27 @@ class MCMCForestSampler {
   ~MCMCForestSampler() {}
 
   void SampleOneIter(ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, 
-                     ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector<double> variable_weights, double global_variance) {
-    
+                     ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector<double> variable_weights, 
+                     double global_variance, bool pre_initialized = false) {
     // Previous number of samples
     int prev_num_samples = forests.NumSamples();
-    // Add new forest to the container
-    forests.AddSamples(1);
 
-    if (prev_num_samples == 0) {
+    if ((prev_num_samples == 0) && (!pre_initialized)) {
+      // Add new forest to the container
+      forests.AddSamples(1);
+      
       // Set initial value for each leaf in the forest
       double root_pred = ComputeMeanOutcome(residual) / static_cast<double>(forests.NumTrees());
       TreeEnsemble* ensemble = forests.GetEnsemble(0);
       leaf_model.SetEnsembleRootPredictedValue(dataset, ensemble, root_pred);
-    } else {
+    } else if (prev_num_samples > 0) {
+      // Add new forest to the container
+      forests.AddSamples(1);
+      
       // Copy previous forest
       forests.CopyFromPreviousSample(prev_num_samples, prev_num_samples - 1);
+    } else {
+      forests.IncrementSampleCount();
     }
 
     // Run the MCMC algorithm for each tree
@@ -462,23 +493,29 @@ class GFRForestSampler {
 
   void SampleOneIter(ForestTracker& tracker, ForestContainer& forests, LeafModel& leaf_model, ForestDataset& dataset, 
                      ColumnVector& residual, TreePrior& tree_prior, std::mt19937& gen, std::vector<double> variable_weights, 
-                     double global_variance, std::vector<FeatureType>& feature_types) {
+                     double global_variance, std::vector<FeatureType>& feature_types, bool pre_initialized = false) {
     // Previous number of samples
     int prev_num_samples = forests.NumSamples();
-    // Add new forest to the container
-    forests.AddSamples(1);
 
-    if (prev_num_samples == 0) {
+    if ((prev_num_samples == 0) && (!pre_initialized)) {
+      // Add new forest to the container
+      forests.AddSamples(1);
+      
       // Set initial value for each leaf in the forest
       double root_pred = ComputeMeanOutcome(residual) / static_cast<double>(forests.NumTrees());
       TreeEnsemble* ensemble = forests.GetEnsemble(0);
       leaf_model.SetEnsembleRootPredictedValue(dataset, ensemble, root_pred);
-    } else {
+    } else if (prev_num_samples > 0) {
+      // Add new forest to the container
+      forests.AddSamples(1);
+
       // NOTE: only doing this for the simplicity of the partial residual step
       // We could alternatively "reach back" to the tree predictions from a previous
       // sample (whenever there is more than one sample). This is cleaner / quicker
       // to implement during this refactor.
       forests.CopyFromPreviousSample(prev_num_samples, prev_num_samples - 1);
+    } else {
+      forests.IncrementSampleCount();
     }
 
     // Run the GFR algorithm for each tree

src/container.cpp

@@ -23,12 +23,37 @@ ForestContainer::ForestContainer(int num_samples, int num_trees, int output_dime
   num_trees_ = num_trees;
   output_dimension_ = output_dimension;
   is_leaf_constant_ = is_leaf_constant;
+  initialized_ = true;
 }
 
 void ForestContainer::CopyFromPreviousSample(int new_sample_id, int previous_sample_id) {
   forests_[new_sample_id].reset(new TreeEnsemble(*forests_[previous_sample_id]));
 }
 
+void ForestContainer::InitializeRoot(double leaf_value) {
+  CHECK(initialized_);
+  CHECK_EQ(num_samples_, 0);
+  CHECK_EQ(forests_.size(), 0);
+  forests_.resize(1);
+  forests_[0].reset(new TreeEnsemble(num_trees_, output_dimension_, is_leaf_constant_));
+  // NOTE: not setting num_samples = 1, since we are just initializing constant root 
+  // nodes and the forest still needs to be sampled by either MCMC or GFR
+  num_samples_ = 0;
+  SetLeafValue(0, leaf_value);
+}
+
+void ForestContainer::InitializeRoot(std::vector<double>& leaf_vector) {
+  CHECK(initialized_);
+  CHECK_EQ(num_samples_, 0);
+  CHECK_EQ(forests_.size(), 0);
+  forests_.resize(1);
+  forests_[0].reset(new TreeEnsemble(num_trees_, output_dimension_, is_leaf_constant_));
+  // NOTE: not setting num_samples = 1, since we are just initializing constant root 
+  // nodes and the forest still needs to be sampled by either MCMC or GFR
+  num_samples_ = 0;
+  SetLeafVector(0, leaf_vector);
+}
+
 void ForestContainer::AddSamples(int num_samples) {
   CHECK(initialized_);
   int total_new_samples = num_samples + num_samples_;

src/leaf_model.cpp

@@ -691,9 +691,10 @@ void GaussianMultivariateRegressionLeafModel::SetEnsembleRootPredictedValue(Fore
     Log::Fatal("For multivariate leaf regression, outcomes should be centered / scaled so that the root coefficients can be initialized to 0");
   }
 
+  std::vector<double> root_pred_vector(ensemble->OutputDimension(), root_pred_value);
   for (int i = 0; i < num_trees; i++) {
     Tree* tree = ensemble->GetTree(i);
-    tree->SetLeaf(0, root_pred_value);
+    tree->SetLeafVector(0, root_pred_vector);
   }
 }
 

src/tree.cpp

@@ -19,29 +19,11 @@ constexpr std::int32_t Tree::kDeletedNodeMarker;
 constexpr std::int32_t Tree::kRoot;
 
 std::int32_t Tree::NumLeaves() const {
-  std::int32_t leaves { 0 };
-  auto const& self = *this;
-  this->WalkTree([&leaves, &self](std::int32_t nidx) {
-                   if (self.IsLeaf(nidx)) {
-                     leaves++;
-                   }
-                   return true;
-                 });
-  return leaves;
+  return leaves_.size();
 }
 
 std::int32_t Tree::NumLeafParents() const {
-  std::int32_t leaf_parents { 0 };
-  auto const& self = *this;
-  this->WalkTree([&leaf_parents, &self](std::int32_t nidx) {
-                   if (!self.IsLeaf(nidx)){
-                     if ((self.IsLeaf(self.LeftChild(nidx))) && (self.IsLeaf(self.RightChild(nidx)))){
-                      leaf_parents++;
-                     }
-                   }
-                   return true;
-                 });
-  return leaf_parents;
+  return leaf_parents_.size();
 }
 
 std::int32_t Tree::NumSplitNodes() const {