Updated kernel code

Author: andrewherren

R/kernel.R

@@ -206,7 +206,6 @@ computeForestLeafVariances <- function(model_object, forest_type, forest_inds=NU
 #' Compute and return the largest possible leaf index computable by `computeForestLeafIndices` for the forests in a designated forest sample container.
 #'
 #' @param model_object Object of type `bartmodel`, `bcfmodel`, or `ForestSamples` corresponding to a BART / BCF model with at least one forest sample, or a low-level `ForestSamples` object.
-#' @param covariates Covariates to use for prediction. Must have the same dimensions / column types as the data used to train a forest.
 #' @param forest_type Which forest to use from `model_object`. 
 #' Valid inputs depend on the model type, and whether or not a 
 #' 
@@ -238,7 +237,7 @@ computeForestLeafVariances <- function(model_object, forest_type, forest_inds=NU
 #' computeForestMaxLeafIndex(bart_model, X, "mean")
 #' computeForestMaxLeafIndex(bart_model, X, "mean", 0)
 #' computeForestMaxLeafIndex(bart_model, X, "mean", c(1,3,9))
-computeForestMaxLeafIndex <- function(model_object, covariates, forest_type=NULL, forest_inds=NULL) {
+computeForestMaxLeafIndex <- function(model_object, forest_type=NULL, forest_inds=NULL) {
     # Extract relevant forest container
     stopifnot(any(c(inherits(model_object, "bartmodel"), inherits(model_object, "bcfmodel"), inherits(model_object, "ForestSamples"))))
     model_type <- ifelse(inherits(model_object, "bartmodel"), "bart", ifelse(inherits(model_object, "bcfmodel"), "bcf", "forest_samples"))

src/kernel.cpp

@@ -11,7 +11,7 @@ typedef Eigen::Map<Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic, Eigen::Col
 
 [[cpp11::register]]
 int forest_container_get_max_leaf_index_cpp(cpp11::external_pointer<StochTree::ForestContainer> forest_container, int forest_num) {
-    return forest_container->GetEnsemble(forest_num)->GetMaxLeafIndex();
+    return forest_container->GetEnsemble(forest_num)->GetMaxLeafIndex() - 1;
 }
 
 [[cpp11::register]]

src/py_stochtree.cpp

@@ -1861,7 +1861,7 @@ py::array_t<int> cppComputeForestContainerLeafIndices(ForestContainerCpp& forest
 }
 
 int cppComputeForestMaxLeafIndex(ForestContainerCpp& forest_container, int forest_num) {
-  return forest_container.GetForest(forest_num)->GetMaxLeafIndex();
+  return forest_container.GetForest(forest_num)->GetMaxLeafIndex() - 1;
 }
 
 void ForestContainerCpp::LoadFromJson(JsonCpp& json, std::string forest_label) {

stochtree/__init__.py

@@ -4,19 +4,22 @@
 from .config import ForestModelConfig, GlobalModelConfig
 from .data import Dataset, Residual
 from .forest import Forest, ForestContainer
-from .kernel import compute_forest_leaf_indices
+from .kernel import (
+    compute_forest_leaf_indices, 
+    compute_forest_max_leaf_index
+)
 from .preprocessing import CovariatePreprocessor
 from .random_effects import (
-  RandomEffectsContainer, 
-  RandomEffectsDataset, 
-  RandomEffectsModel, 
-  RandomEffectsTracker, 
+    RandomEffectsContainer, 
+    RandomEffectsDataset, 
+    RandomEffectsModel, 
+    RandomEffectsTracker, 
 )
 from .sampler import (
-  RNG, 
-  ForestSampler, 
-  GlobalVarianceModel, 
-  LeafVarianceModel
+    RNG, 
+    ForestSampler, 
+    GlobalVarianceModel, 
+    LeafVarianceModel
 )
 from .serialization import JSONSerializer
 from .utils import (
@@ -58,5 +61,6 @@
     "_standardize_array_to_list",
     "_standardize_array_to_np",
     "compute_forest_leaf_indices",
+    "compute_forest_max_leaf_index", 
     "calibrate_global_error_variance",
 ]

stochtree/kernel.py

@@ -93,8 +93,109 @@ def compute_forest_leaf_indices(model_object: Union[BARTModel, BCFModel, ForestC
     num_forests = forest_container.num_samples()
     if forest_inds is None:
         forest_inds = np.arange(num_forests)
-    else:
+    elif isinstance(forest_inds, int):
+        if not forest_inds >= 0 or not forest_inds < num_forests:
+            raise ValueError("The index in forest_inds must be >= 0 and < the total number of samples in a forest container")
+        forest_inds = np.array([forest_inds])
+    elif isinstance(forest_inds, np.ndarray):
+        if forest_inds.size > 1:
+            forest_inds = np.squeeze(forest_inds)
+            if forest_inds.ndim > 1:
+                raise ValueError("forest_inds must be a one-dimensional numpy array")
         if not np.all(forest_inds >= 0) or not np.all(forest_inds < num_forests):
             raise ValueError("The indices in forest_inds must be >= 0 and < the total number of samples in a forest container")
+    else:
+        raise ValueError("forest_inds must be a one-dimensional numpy array")
 
     return cppComputeForestContainerLeafIndices(forest_container.forest_container_cpp, covariates_processed, forest_inds)
+
+def compute_forest_max_leaf_index(model_object: Union[BARTModel, BCFModel, ForestContainer], forest_type: str = None, forest_inds: Union[int, np.ndarray] = None):
+    """
+    Compute and return the largest possible leaf index computable by `compute_forest_leaf_indices` for the forests in a designated forest sample container.
+
+    Parameters
+    ----------
+    model_object : BARTModel, BCFModel, or ForestContainer
+        Object corresponding to a BART / BCF model with at least one forest sample, or a low-level `ForestContainer` object.
+    forest_type : str
+        Which forest to use from `model_object`. Valid inputs depend on the model type, and whether or not a given forest was sampled in that model.
+
+            * **BART**
+                * `'mean'`: `'mean'`: Extracts leaf indices for the mean forest
+                * `'variance'`: Extracts leaf indices for the variance forest
+            * **BCF**
+                * `'prognostic'`: Extracts leaf indices for the prognostic forest
+                * `'treatment'`: Extracts leaf indices for the treatment effect forest
+                * `'variance'`: Extracts leaf indices for the variance forest
+            * **ForestContainer**
+                * `NULL`: It is not necessary to disambiguate when this function is called directly on a `ForestSamples` object. This is the default value of this
+    
+    forest_inds : int or np.ndarray
+        Indices of the forest sample(s) for which to compute max leaf indices. If not provided, this function will return max leaf indices for every sample of a forest. 
+        This function uses 0-indexing, so the first forest sample corresponds to `forest_num = 0`, and so on.
+    
+    Returns
+    -------
+    Numpy array containing the largest possible leaf index computable by `compute_forest_leaf_indices` for the forests in a designated forest sample container.
+    """
+    # Extract relevant forest container
+    if not isinstance(model_object, BARTModel) and not isinstance(model_object, BCFModel) and not isinstance(model_object, ForestContainer):
+        raise ValueError("model_object must be one of BARTModel, BCFModel, or ForestContainer")
+    if isinstance(model_object, BARTModel):
+        model_type = "bart"
+        if forest_type is None:
+            raise ValueError("forest_type must be specified for a BARTModel model_type (either set to 'mean' or 'variance')")
+    elif isinstance(model_object, BCFModel):
+        model_type = "bcf"
+        if forest_type is None:
+            raise ValueError("forest_type must be specified for a BCFModel model_type (either set to 'prognostic', 'treatment' or 'variance')")
+    else:
+        model_type = "forest"
+    if model_type == "bart":
+        if forest_type == "mean":
+            if not model_object.include_mean_forest:
+                raise ValueError("Mean forest was not sampled for model_object, but requested by forest_type")
+            forest_container = model_object.forest_container_mean
+        else:
+            if not model_object.include_variance_forest:
+                raise ValueError("Variance forest was not sampled for model_object, but requested by forest_type")
+            forest_container = model_object.forest_container_variance
+    elif model_type == "bcf":
+        if forest_type == "prognostic":
+            forest_container = model_object.forest_container_mu
+        elif forest_type == "treatment":
+            forest_container = model_object.forest_container_tau
+        else:
+            if not model_object.include_variance_forest:
+                raise ValueError("Variance forest was not sampled for model_object, but requested by forest_type")
+            forest_container = model_object.forest_container_variance
+    else:
+        forest_container = model_object
+    
+    # Preprocess forest indices
+    num_forests = forest_container.num_samples()
+    if forest_inds is None:
+        forest_inds = np.arange(num_forests)
+    elif isinstance(forest_inds, int):
+        if not forest_inds >= 0 or not forest_inds < num_forests:
+            raise ValueError("The index in forest_inds must be >= 0 and < the total number of samples in a forest container")
+        forest_inds = np.array([forest_inds])
+    elif isinstance(forest_inds, np.ndarray):
+        if forest_inds.size > 1:
+            forest_inds = np.squeeze(forest_inds)
+            if forest_inds.ndim > 1:
+                raise ValueError("forest_inds must be a one-dimensional numpy array")
+        if not np.all(forest_inds >= 0) or not np.all(forest_inds < num_forests):
+            raise ValueError("The indices in forest_inds must be >= 0 and < the total number of samples in a forest container")
+    else:
+        raise ValueError("forest_inds must be a one-dimensional numpy array")
+    
+    # Compute max index
+    output_size = len(forest_inds)
+    output = np.empty(output_size)
+    for i in np.arange(output_size):
+        output[i] = cppComputeForestMaxLeafIndex(forest_container.forest_container_cpp, forest_inds[i])
+    
+    # Return result    
+    return output
+

test/python/test_kernel.py

@@ -5,7 +5,8 @@
     Dataset,
     Forest,
     ForestContainer,
-    compute_forest_leaf_indices
+    compute_forest_leaf_indices, 
+    compute_forest_max_leaf_index
 )
 
 
@@ -35,6 +36,7 @@ def test_forest(self):
 
         # Check that regular and "raw" predictions are the same (since the leaf is constant)
         computed = compute_forest_leaf_indices(forest_samples, X)
+        max_leaf_index = compute_forest_max_leaf_index(forest_samples)
         expected = np.array([
             [0], 
             [0], 
@@ -52,12 +54,14 @@ def test_forest(self):
 
         # Assertion
         np.testing.assert_almost_equal(computed, expected)
+        assert max_leaf_index == [2]
 
         # Split the left leaf of the first tree in the ensemble at X[,2] > 4.0
         forest_samples.add_numeric_split(0, 0, 1, 1, 4.0, -7.5, -2.5)
 
         # Check that regular and "raw" predictions are the same (since the leaf is constant)
         computed = compute_forest_leaf_indices(forest_samples, X)
+        max_leaf_index = compute_forest_max_leaf_index(forest_samples)
         expected = np.array([
             [2], 
             [1], 
@@ -75,3 +79,4 @@ def test_forest(self):
 
         # Assertion
         np.testing.assert_almost_equal(computed, expected)
+        assert max_leaf_index == [3]