Added forest container methods to python interface as well

Author: andrewherren

R/forest.R

@@ -36,7 +36,7 @@ ForestSamples <- R6::R6Class(
 
         #' @description
         #' Add a constant value to every leaf of every tree of a given forest
-        #' @param forest_index Index of forest whos leaves will be modified (0-indexed)
+        #' @param forest_index Index of forest whose leaves will be modified (0-indexed)
         #' @param constant_value Value to add to every leaf of every tree of the forest at `forest_index`
         add_to_forest = function(forest_index, constant_value) {
             stopifnot(forest_index < self$num_samples())
@@ -46,7 +46,7 @@ ForestSamples <- R6::R6Class(
 
         #' @description
         #' Multiply every leaf of every tree of a given forest by constant value
-        #' @param forest_index Index of forest whos leaves will be modified (0-indexed)
+        #' @param forest_index Index of forest whose leaves will be modified (0-indexed)
         #' @param constant_multiple Value to multiply through by every leaf of every tree of the forest at `forest_index`
         multiply_forest = function(forest_index, constant_multiple) {
             stopifnot(forest_index < self$num_samples())

include/stochtree/container.h

@@ -59,7 +59,7 @@ class ForestContainer {
   /*!
    * \brief Add a constant value to every leaf of every tree of a specified forest
    * 
-   * \param forest_index Index of forest whos leaves will be modified
+   * \param forest_index Index of forest whose leaves will be modified
    * \param constant_value Value to add to every leaf of every tree of the forest at `forest_index`
    */
   void AddToForest(int forest_index, double constant_value) {
@@ -68,7 +68,7 @@ class ForestContainer {
   /*!
    * \brief Multiply every leaf of every tree of a specified forest by a constant value
    * 
-   * \param forest_index Index of forest whos leaves will be modified
+   * \param forest_index Index of forest whose leaves will be modified
    * \param constant_multiple Value to multiply through by every leaf of every tree of the forest at `forest_index`
    */
   void MultiplyForest(int forest_index, double constant_multiple) {

src/py_stochtree.cpp

@@ -173,6 +173,21 @@ class ForestContainerCpp {
   }
   ~ForestContainerCpp() {}
 
+  void CombineForests(py::array_t<int> forest_inds) {
+    int num_forests = forest_inds.size();
+    for (int j = 1; j < num_forests; j++) {
+      forest_samples_->MergeForests(forest_inds.at(0), forest_inds.at(j));
+    }
+  }
+
+  void AddToForest(int forest_index, double constant_value) {
+    forest_samples_->AddToForest(forest_index, constant_value);
+  }
+
+  void MultiplyForest(int forest_index, double constant_multiple) {
+    forest_samples_->MultiplyForest(forest_index, constant_multiple);
+  }
+
   int OutputDimension() {
     return forest_samples_->OutputDimension();
   }
@@ -2023,9 +2038,12 @@ PYBIND11_MODULE(stochtree_cpp, m) {
 
   py::class_<RngCpp>(m, "RngCpp")
     .def(py::init<int>());
-
+  
   py::class_<ForestContainerCpp>(m, "ForestContainerCpp")
     .def(py::init<int,int,bool,bool>())
+    .def("CombineForests", &ForestContainerCpp::CombineForests)
+    .def("AddToForest", &ForestContainerCpp::AddToForest)
+    .def("MultiplyForest", &ForestContainerCpp::MultiplyForest)
     .def("OutputDimension", &ForestContainerCpp::OutputDimension)
     .def("NumTrees", &ForestContainerCpp::NumTrees)
     .def("NumSamples", &ForestContainerCpp::NumSamples)

stochtree/forest.py

@@ -159,6 +159,63 @@ def set_root_leaves(
         else:
             self.forest_container_cpp.SetRootValue(forest_num, leaf_value)
 
+    def combine_forests(
+        self, forest_inds: np.array
+    ) -> None:
+        """
+        Collapse specified forests into a single forest
+
+        Parameters
+        ----------
+        forest_inds : np.array
+            Indices of forests to be combined (0-indexed).
+        """
+        if not isinstance(forest_inds, np.ndarray):
+            raise ValueError("forest_inds must be either a np.array")
+        if not np.issubdtype(forest_inds.dtype, np.number):
+            raise ValueError("forest_inds must be an integer-convertible np.array")
+        forest_inds_sorted = np.sort(forest_inds)
+        forest_inds_sorted = np.astype(forest_inds_sorted, int)
+        self.forest_container_cpp.CombineForests(forest_inds_sorted)
+
+    def add_to_forest(
+        self, forest_index: int, constant_value : float
+    ) -> None:
+        """
+        Add a constant value to every leaf of every tree of a given forest
+
+        Parameters
+        ----------
+        forest_index : int 
+            Index of forest whose leaves will be modified (0-indexed)
+        constant_value : float 
+            Value to add to every leaf of every tree of the forest at `forest_index`
+        """
+        if not isinstance(forest_index, int) and not isinstance(constant_value, (int, float)):
+            raise ValueError("forest_index must be an integer and constant_multiple must be a float or int")
+        if not forest_index >= 0 or not forest_index < self.forest_container_cpp.NumSamples():
+            raise ValueError("forest_index must be >= 0 and less than the total number of samples in a forest container")
+        self.forest_container_cpp.AddToForest(forest_index, constant_value)
+
+    def multiply_forest(
+        self, forest_index: int, constant_multiple : float
+    ) -> None:
+        """
+        Multiply every leaf of every tree of a given forest by constant value
+
+        Parameters
+        ----------
+        forest_index : int 
+            Index of forest whose leaves will be modified (0-indexed)
+        constant_multiple : float 
+            Value to multiply through by every leaf of every tree of the forest at `forest_index`
+        """
+        if not isinstance(forest_index, int) and not isinstance(constant_multiple, (int, float)):
+            raise ValueError("forest_index must be an integer and constant_multiple must be a float or int")
+        if not forest_index >= 0 or not forest_index < self.forest_container_cpp.NumSamples():
+            raise ValueError("forest_index must be >= 0 and less than the total number of samples in a forest container")
+        self.forest_container_cpp.MultiplyForest(forest_index, constant_multiple)
+
     def save_to_json_file(self, json_filename: str) -> None:
         """
         Save the forests in the container to a JSON file.

test/R/testthat/test-forest-container.R

@@ -57,7 +57,7 @@ test_that("Univariate constant forest container", {
     # Add 1.0 to every tree in first forest
     forest_samples$add_to_forest(0, 1.0)
 
-    # Check that predictions are + num_trees
+    # Check that predictions are += num_trees
     pred_expected <- pred + num_trees
     pred <- forest_samples$predict(forest_dataset)
 

test/python/test_forest_container.py

@@ -0,0 +1,77 @@
+import numpy as np
+
+from stochtree import Dataset, ForestContainer
+
+
+class TestPredict:
+    def test_constant_leaf_prediction(self):
+        # Create dataset
+        X = np.array(
+            [[1.5, 8.7, 1.2],
+             [2.7, 3.4, 5.4],
+             [3.6, 1.2, 9.3],
+             [4.4, 5.4, 10.4],
+             [5.3, 9.3, 3.6],
+             [6.1, 10.4, 4.4]]
+        )
+        n, p = X.shape
+        num_trees = 10
+        output_dim = 1
+        forest_dataset = Dataset()
+        forest_dataset.add_covariates(X)
+        forest_samples = ForestContainer(num_trees, output_dim, True, False)
+
+        # Initialize a forest with constant root predictions
+        forest_samples.add_sample(0.)
+
+        # Split the root of the first tree in the ensemble at X[,1] > 4.0
+        # and then split the left leaf of the first tree in the ensemble at X[,2] > 4.0
+        forest_samples.add_numeric_split(0, 0, 0, 0, 4.0, -5., 5.)
+        forest_samples.add_numeric_split(0, 0, 1, 1, 4.0, -7.5, -2.5)
+
+        # Store the predictions of the "original" forest before modifications
+        pred_orig = forest_samples.predict(forest_dataset)
+
+        # Multiply first forest by 2.0
+        forest_samples.multiply_forest(0, 2.0)
+        
+        # Check that predictions are all double
+        pred = forest_samples.predict(forest_dataset)
+        pred_expected = pred_orig * 2.0
+        
+        # Assertion
+        np.testing.assert_almost_equal(pred, pred_expected)
+
+        # Add 1.0 to every tree in first forest
+        forest_samples.add_to_forest(0, 1.0)
+        
+        # Check that predictions are += num_trees
+        pred_expected = pred + num_trees
+        pred = forest_samples.predict(forest_dataset)
+        
+        # Assertion
+        np.testing.assert_almost_equal(pred, pred_expected)
+
+        # Initialize a new forest with constant root predictions
+        forest_samples.add_sample(0.)
+
+        # Split the second forest as the first forest was split
+        forest_samples.add_numeric_split(1, 0, 0, 0, 4.0, -5., 5.)
+        forest_samples.add_numeric_split(1, 0, 1, 1, 4.0, -7.5, -2.5)
+        
+        # Check that predictions are as expected
+        pred_expected_new = np.c_[pred_expected, pred_orig]
+        pred = forest_samples.predict(forest_dataset)
+        
+        # Assertion
+        np.testing.assert_almost_equal(pred, pred_expected_new)
+
+        # Combine second forest with the first forest
+        forest_samples.combine_forests(np.array([0,1]))
+
+        # Check that predictions are as expected
+        pred_expected_new = np.c_[pred_expected + pred_orig, pred_orig]
+        pred = forest_samples.predict(forest_dataset)
+        
+        # Assertion
+        np.testing.assert_almost_equal(pred, pred_expected_new)