Merge pull request #100 from StochasticTree/forest_leaf_scale_enhancement

Added function to extract forest leaf scale parameters

Author: andrewherren

NAMESPACE

@@ -8,6 +8,7 @@ export(bart)
 export(bcf)
 export(calibrate_inverse_gamma_error_variance)
 export(computeForestLeafIndices)
+export(computeForestLeafVariances)
 export(computeMaxLeafIndex)
 export(convertBARTModelToJson)
 export(convertBCFModelToJson)

R/kernel.R

@@ -1,7 +1,10 @@
-#' Compute and return a vector representation of a forest's leaf predictions for 
+#' Compute vector of forest leaf indices
+#' 
+#' @description Compute and return a vector representation of a forest's leaf predictions for 
 #' every observation in a dataset.
-#' The vector has a "column-major" format that can be easily re-represented as 
-#' as a CSC sparse matrix: elements are organized so that the first `n` elements 
+#' 
+#' The vector has a "row-major" format that can be easily re-represented as 
+#' as a CSR sparse matrix: elements are organized so that the first `n` elements 
 #' correspond to leaf predictions for all `n` observations in a dataset for the 
 #' first tree in an ensemble, the next `n` elements correspond to predictions for 
 #' the second tree and so on. The "data" for each element corresponds to a uniquely 
@@ -12,7 +15,7 @@
 #' @param model_object Object of type `bartmodel` or `bcf` corresponding to a BART / BCF model with at least one forest sample
 #' @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 
+#' Valid inputs depend on the model type, and whether or not a given forest was sampled in that model.
 #' 
 #'   **1. BART**
 #'
@@ -88,6 +91,89 @@ computeForestLeafIndices <- function(model_object, covariates, forest_type, fore
     return(leaf_ind_matrix)
 }
 
+#' Compute vector of forest leaf scale parameters
+#' 
+#' @description Return each forest's leaf node scale parameters.
+#' 
+#' If leaf scale is not sampled for the forest in question, throws an error that the 
+#' leaf model does not have a stochastic scale parameter.
+#' 
+#' @param model_object Object of type `bartmodel` or `bcf` corresponding to a BART / BCF model with at least one forest sample
+#' @param forest_type 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.
+#' 
+#'   **1. BART**
+#'
+#'   - `'mean'`: Extracts leaf indices for the mean forest
+#'   - `'variance'`: Extracts leaf indices for the variance forest
+#' 
+#'   **2. 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
+#' 
+#' @param forest_inds (Optional) Indices of the forest sample(s) for which to compute leaf indices. If not provided, 
+#' this function will return leaf indices for every sample of a forest. 
+#' This function uses 1-indexing, so the first forest sample corresponds to `forest_num = 1`, and so on.
+#' @return Vector of size `length(forest_inds)` with the leaf scale parameter for each requested forest.
+#' @export
+computeForestLeafVariances <- function(model_object, forest_type, forest_inds=NULL) {
+    # Extract relevant forest container
+    stopifnot(class(model_object) %in% c("bartmodel", "bcf"))
+    model_type <- ifelse(class(model_object)=="bartmodel", "bart", "bcf")
+    if (model_type == "bart") {
+        stopifnot(forest_type %in% c("mean", "variance"))
+        if (forest_type=="mean") {
+            if (!model_object$model_params$include_mean_forest) {
+                stop("Mean forest was not sampled in the bart model provided")
+            }
+            if (model_object$model_params$sample_sigma_leaf == F) {
+                stop("Leaf scale parameter was not sampled for the mean forest in the bart model provided")
+            }
+            leaf_scale_vector <- model_object$sigma2_leaf_samples
+        } else if (forest_type=="variance") {
+            if (!model_object$model_params$include_variance_forest) {
+                stop("Variance forest was not sampled in the bart model provided")
+            }
+            stop("Leaf scale parameter was not sampled for the variance forest in the bart model provided")
+        }
+    } else {
+        stopifnot(forest_type %in% c("prognostic", "treatment", "variance"))
+        if (forest_type=="prognostic") {
+            if (model_object$model_params$sample_sigma_leaf_mu == F) {
+                stop("Leaf scale parameter was not sampled for the prognostic forest in the bcf model provided")
+            }
+            leaf_scale_vector <- model_object$sigma_leaf_mu_samples
+        } else if (forest_type=="treatment") {
+            if (model_object$model_params$sample_sigma_leaf_tau == F) {
+                stop("Leaf scale parameter was not sampled for the treatment effect forest in the bcf model provided")
+            }
+            leaf_scale_vector <- model_object$sigma_leaf_tau_samples
+        } else if (forest_type=="variance") {
+            if (!model_object$model_params$include_variance_forest) {
+                stop("Variance forest was not sampled in the bcf model provided")
+            }
+            stop("Leaf scale parameter was not sampled for the variance forest in the bcf model provided")
+        }
+    }
+    
+    # Preprocess forest indices
+    num_forests <- forest_container$num_samples()
+    if (is.null(forest_inds)) {
+        forest_inds <- as.integer(1:num_forests)
+    } else {
+        stopifnot(all(forest_inds <= num_forests))
+        stopifnot(all(forest_inds >= 1))
+        forest_inds <- as.integer(forest_inds)
+    }
+    
+    # Gather leaf scale parameters
+    leaf_scale_params <- leaf_scale_vector[forest_inds]
+
+    return(leaf_scale_params)
+}
+
 #' 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` or `bcf` corresponding to a BART / BCF model with at least one forest sample

_pkgdown.yml

@@ -86,8 +86,9 @@ reference:
   - calibrate_inverse_gamma_error_variance
   - preprocessBartParams
   - preprocessBcfParams
-  - computeForestLeafIndices
   - computeMaxLeafIndex
+  - computeForestLeafIndices
+  - computeForestLeafVariances
 
 - subtitle: Random Effects
   desc: >
@@ -104,7 +105,6 @@ reference:
   - getRandomEffectSamples.bcf
   - sample_sigma2_one_iteration
   - sample_tau_one_iteration
-  - sample_tau_one_iteration
 
 - title: Package info
   desc: >