Enable access to forest leaf values and counts at high level

R/cpp11.R

@@ -252,6 +252,14 @@ average_max_depth_forest_container_cpp <- function(forest_samples) {
   .Call(`_stochtree_average_max_depth_forest_container_cpp`, forest_samples)
 }
 
+num_leaves_ensemble_forest_container_cpp <- function(forest_samples, forest_num) {
+  .Call(`_stochtree_num_leaves_ensemble_forest_container_cpp`, forest_samples, forest_num)
+}
+
+sum_leaves_squared_ensemble_forest_container_cpp <- function(forest_samples, forest_num) {
+  .Call(`_stochtree_sum_leaves_squared_ensemble_forest_container_cpp`, forest_samples, forest_num)
+}
+
 num_trees_forest_container_cpp <- function(forest_samples) {
   .Call(`_stochtree_num_trees_forest_container_cpp`, forest_samples)
 }

R/forest.R

@@ -323,6 +323,22 @@ ForestSamples <- R6::R6Class(
         #' @return Average maximum depth
         average_max_depth = function() {
             return(average_max_depth_forest_container_cpp(self$forest_container_ptr))
+        }, 
+
+        #' @description
+        #' Number of leaves in a given ensemble in a `ForestContainer` object
+        #' @param forest_num Index of the ensemble to be queried
+        #' @return Count of leaves in the ensemble stored at `forest_num`
+        num_leaves = function(forest_num) {
+            return(num_leaves_ensemble_forest_container_cpp(self$forest_container_ptr, forest_num))
+        }, 
+
+        #' @description
+        #' Sum of squared (raw) leaf values in a given ensemble in a `ForestContainer` object
+        #' @param forest_num Index of the ensemble to be queried
+        #' @return Average maximum depth
+        sum_leaves_squared = function(forest_num) {
+            return(sum_leaves_squared_ensemble_forest_container_cpp(self$forest_container_ptr, forest_num))
         }
     )
 )

src/cpp11.cpp

@@ -469,6 +469,20 @@ extern "C" SEXP _stochtree_average_max_depth_forest_container_cpp(SEXP forest_sa
   END_CPP11
 }
 // forest.cpp
+int num_leaves_ensemble_forest_container_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples, int forest_num);
+extern "C" SEXP _stochtree_num_leaves_ensemble_forest_container_cpp(SEXP forest_samples, SEXP forest_num) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(num_leaves_ensemble_forest_container_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::external_pointer<StochTree::ForestContainer>>>(forest_samples), cpp11::as_cpp<cpp11::decay_t<int>>(forest_num)));
+  END_CPP11
+}
+// forest.cpp
+double sum_leaves_squared_ensemble_forest_container_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples, int forest_num);
+extern "C" SEXP _stochtree_sum_leaves_squared_ensemble_forest_container_cpp(SEXP forest_samples, SEXP forest_num) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(sum_leaves_squared_ensemble_forest_container_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::external_pointer<StochTree::ForestContainer>>>(forest_samples), cpp11::as_cpp<cpp11::decay_t<int>>(forest_num)));
+  END_CPP11
+}
+// forest.cpp
 int num_trees_forest_container_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples);
 extern "C" SEXP _stochtree_num_trees_forest_container_cpp(SEXP forest_samples) {
   BEGIN_CPP11
@@ -1034,6 +1048,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_stochtree_json_load_string_cpp",                                (DL_FUNC) &_stochtree_json_load_string_cpp,                                 2},
     {"_stochtree_json_save_file_cpp",                                  (DL_FUNC) &_stochtree_json_save_file_cpp,                                   2},
     {"_stochtree_json_save_forest_container_cpp",                      (DL_FUNC) &_stochtree_json_save_forest_container_cpp,                       2},
+    {"_stochtree_num_leaves_ensemble_forest_container_cpp",            (DL_FUNC) &_stochtree_num_leaves_ensemble_forest_container_cpp,             2},
     {"_stochtree_num_samples_forest_container_cpp",                    (DL_FUNC) &_stochtree_num_samples_forest_container_cpp,                     1},
     {"_stochtree_num_trees_forest_container_cpp",                      (DL_FUNC) &_stochtree_num_trees_forest_container_cpp,                       1},
     {"_stochtree_output_dimension_forest_container_cpp",               (DL_FUNC) &_stochtree_output_dimension_forest_container_cpp,                1},
@@ -1089,6 +1104,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_stochtree_set_leaf_value_forest_container_cpp",                 (DL_FUNC) &_stochtree_set_leaf_value_forest_container_cpp,                  2},
     {"_stochtree_set_leaf_vector_forest_container_cpp",                (DL_FUNC) &_stochtree_set_leaf_vector_forest_container_cpp,                 2},
     {"_stochtree_subtract_from_column_vector_cpp",                     (DL_FUNC) &_stochtree_subtract_from_column_vector_cpp,                      2},
+    {"_stochtree_sum_leaves_squared_ensemble_forest_container_cpp",    (DL_FUNC) &_stochtree_sum_leaves_squared_ensemble_forest_container_cpp,     2},
     {"_stochtree_tree_prior_cpp",                                      (DL_FUNC) &_stochtree_tree_prior_cpp,                                       4},
     {NULL, NULL, 0}
 };

src/forest.cpp

@@ -96,6 +96,18 @@ double average_max_depth_forest_container_cpp(cpp11::external_pointer<StochTree:
     return forest_samples->AverageMaxDepth();
 }
 
+[[cpp11::register]]
+int num_leaves_ensemble_forest_container_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples, int forest_num) {
+    StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num);
+    return forest->NumLeaves();
+}
+
+[[cpp11::register]]
+double sum_leaves_squared_ensemble_forest_container_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples, int forest_num) {
+    StochTree::TreeEnsemble* forest = forest_samples->GetEnsemble(forest_num);
+    return forest->SumLeafSquared();
+}
+
 [[cpp11::register]]
 int num_trees_forest_container_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_samples) {
     return forest_samples->NumTrees();

src/py_stochtree.cpp

@@ -164,6 +164,16 @@ class ForestContainerCpp {
     return forest_samples_->NumSamples();
   }
 
+  int NumLeaves(int forest_num) {
+    StochTree::TreeEnsemble* forest = forest_samples_->GetEnsemble(forest_num);
+    return forest->NumLeaves();
+  }
+
+  double SumLeafSquared(int forest_num) {
+    StochTree::TreeEnsemble* forest = forest_samples_->GetEnsemble(forest_num);
+    return forest->SumLeafSquared();
+  }
+
   py::array_t<double> Predict(ForestDatasetCpp& dataset) {
     // Predict from the forest container
     data_size_t n = dataset.NumRows();
@@ -1028,7 +1038,9 @@ PYBIND11_MODULE(stochtree_cpp, m) {
     .def("GetTreeSplitCounts", &ForestContainerCpp::GetTreeSplitCounts)
     .def("GetForestSplitCounts", &ForestContainerCpp::GetForestSplitCounts)
     .def("GetOverallSplitCounts", &ForestContainerCpp::GetOverallSplitCounts)
-    .def("GetGranularSplitCounts", &ForestContainerCpp::GetGranularSplitCounts);
+    .def("GetGranularSplitCounts", &ForestContainerCpp::GetGranularSplitCounts)
+    .def("NumLeaves", &ForestContainerCpp::NumLeaves)
+    .def("SumLeafSquared", &ForestContainerCpp::SumLeafSquared);
 
   py::class_<ForestSamplerCpp>(m, "ForestSamplerCpp")
     .def(py::init<ForestDatasetCpp&, py::array_t<int>, int, data_size_t, double, double, int, int>())

stochtree/forest.py

@@ -148,4 +148,22 @@ def get_granular_split_counts(self, num_features: int) -> np.array:
             Total number of features in the training set
         """
         return self.forest_container_cpp.GetGranularSplitCounts(num_features)
+
+    def num_leaves(self, forest_num: int) -> int:
+        """
+        Return the total number of leaves for a given forest in the ``ForestContainer``
+
+        forest_num : :obj:`int`
+            Index of the forest to be queried
+        """
+        return self.forest_container_cpp.NumLeaves(forest_num)
+
+    def sum_leaves_squared(self, forest_num: int) -> float:
+        """
+        Return the total sum of squared leaf values for a given forest in the ``ForestContainer``
+
+        forest_num : :obj:`int`
+            Index of the forest to be queried
+        """
+        return self.forest_container_cpp.SumLeafSquared(forest_num)