Merge pull request #105 from StochasticTree/single-tree-predict

Added ability to predict from a single tree for a given forest

Author: andrewherren

include/stochtree/container.h

@@ -33,9 +33,11 @@ class ForestContainer {
   std::vector<double> Predict(ForestDataset& dataset);
   std::vector<double> PredictRaw(ForestDataset& dataset);
   std::vector<double> PredictRaw(ForestDataset& dataset, int forest_num);
+  std::vector<double> PredictRawSingleTree(ForestDataset& dataset, int forest_num, int tree_num);
   void PredictInPlace(ForestDataset& dataset, std::vector<double>& output);
   void PredictRawInPlace(ForestDataset& dataset, std::vector<double>& output);
   void PredictRawInPlace(ForestDataset& dataset, int forest_num, std::vector<double>& output);
+  void PredictRawSingleTreeInPlace(ForestDataset& dataset, int forest_num, int tree_num, std::vector<double>& output);
   void PredictLeafIndicesInplace(Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& covariates, 
                                  Eigen::Map<Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& output, 
                                  std::vector<int>& forest_indices, int num_trees, data_size_t n);

src/container.cpp

@@ -90,6 +90,14 @@ std::vector<double> ForestContainer::PredictRaw(ForestDataset& dataset, int fore
   return output;
 }
 
+std::vector<double> ForestContainer::PredictRawSingleTree(ForestDataset& dataset, int forest_num, int tree_num) {
+  data_size_t n = dataset.NumObservations();
+  data_size_t total_output_size = n * output_dimension_;
+  std::vector<double> output(total_output_size);
+  PredictRawSingleTreeInPlace(dataset, forest_num, tree_num, output);
+  return output;
+}
+
 void ForestContainer::PredictInPlace(ForestDataset& dataset, std::vector<double>& output) {
   data_size_t n = dataset.NumObservations();
   data_size_t total_output_size = n*num_samples_;
@@ -123,6 +131,14 @@ void ForestContainer::PredictRawInPlace(ForestDataset& dataset, int forest_num,
   forests_[forest_num]->PredictRawInplace(dataset, output, 0, num_trees, offset);
 }
 
+void ForestContainer::PredictRawSingleTreeInPlace(ForestDataset& dataset, int forest_num, int tree_num, std::vector<double>& output) {
+  data_size_t n = dataset.NumObservations();
+  data_size_t total_output_size = n * output_dimension_;
+  CHECK_EQ(total_output_size, output.size());
+  data_size_t offset = 0;
+  forests_[forest_num]->PredictRawInplace(dataset, output, tree_num, tree_num+1, offset);
+}
+
 void ForestContainer::PredictLeafIndicesInplace(
   Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& covariates, 
   Eigen::Map<Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& output, 

src/forest.cpp

@@ -422,3 +422,21 @@ cpp11::writable::doubles_matrix<> predict_forest_raw_single_forest_cpp(cpp11::ex
 
     return output;
 }
+
+[[cpp11::register]]
+cpp11::writable::doubles_matrix<> predict_forest_raw_single_tree_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples, cpp11::external_pointer<StochTree::ForestDataset> dataset, int forest_num, int tree_num) {
+    // Predict from the sampled forests
+    std::vector<double> output_raw = forest_samples->PredictRawSingleTree(*dataset, forest_num, tree_num);
+    
+    // Convert result to a matrix
+    int n = dataset->GetCovariates().rows();
+    int output_dimension = forest_samples->OutputDimension();
+    cpp11::writable::doubles_matrix<> output(n, output_dimension);
+    for (size_t i = 0; i < n; i++) {
+        for (int j = 0; j < output_dimension; j++) {
+            output(i, j) = output_raw[i*output_dimension + j];
+        }
+    }
+    
+    return output;
+}

src/py_stochtree.cpp

@@ -235,6 +235,29 @@ class ForestContainerCpp {
     return result;
   }
 
+  py::array_t<double> PredictRawSingleTree(ForestDatasetCpp& dataset, int forest_num, int tree_num) {
+    // Predict from the forest container
+    data_size_t n = dataset.NumRows();
+    int num_samples = this->NumSamples();
+    int output_dim = this->OutputDimension();
+    StochTree::ForestDataset* data_ptr = dataset.GetDataset();
+    std::vector<double> output_raw = forest_samples_->PredictRawSingleTree(*data_ptr, forest_num, tree_num);
+
+    // Convert result to a matrix
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({n, output_dim}));
+    auto accessor = result.mutable_unchecked<2>();
+    // py::buffer_info buf = result.request();
+    // double *ptr = static_cast<double *>(buf.ptr);
+    for (size_t i = 0; i < n; i++) {
+      for (int j = 0; j < output_dim; j++) {
+        accessor(i,j) = output_raw[i*output_dim + j];
+        // ptr[i*output_dim + j] = output_raw[i*output_dim + j];
+      }
+    }
+
+    return result;
+  }
+
   void SetRootValue(int forest_num, double leaf_value) {
     forest_samples_->InitializeRoot(leaf_value);
   }

stochtree/forest.py

@@ -28,6 +28,10 @@ def predict_raw_single_forest(self, dataset: Dataset, forest_num: int) -> np.arr
         # Predict raw leaf values for a specific forest (indexed by forest_num) from Dataset
         return self.forest_container_cpp.PredictRawSingleForest(dataset.dataset_cpp, forest_num)
 
+    def predict_raw_single_tree(self, dataset: Dataset, forest_num: int, tree_num: int) -> np.array:
+        # Predict raw leaf values for a specific tree from specific forest from Dataset
+        return self.forest_container_cpp.PredictRawSingleTree(dataset.dataset_cpp, forest_num, tree_num)
+    
     def set_root_leaves(self, forest_num: int, leaf_value: Union[float, np.array]) -> None:
         # Predict raw leaf values for a specific forest (indexed by forest_num) from Dataset
         if not isinstance(leaf_value, np.ndarray) and not isinstance(leaf_value, float):