Initial BART refactor that avoids double predicting forests on the train dataset

Author: andrewherren

R/bart.R

@@ -707,6 +707,8 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
     num_retained_samples <- num_gfr + ifelse(keep_burnin, num_burnin, 0) + num_mcmc * num_chains
     if (sample_sigma2_global) global_var_samples <- rep(NA, num_retained_samples)
     if (sample_sigma2_leaf) leaf_scale_samples <- rep(NA, num_retained_samples)
+    if (include_mean_forest) mean_forest_pred_train <- matrix(NA_real_, nrow(X_train), num_retained_samples)
+    if (include_variance_forest) variance_forest_pred_train <- matrix(NA_real_, nrow(X_train), num_retained_samples)
     sample_counter <- 0
 
     # Initialize the leaves of each tree in the mean forest
@@ -757,13 +759,23 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                     active_forest = active_forest_mean, rng = rng, forest_model_config = forest_model_config_mean, 
                     global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
                 )
+                
+                # Cache predictions
+                if (keep_sample) {
+                    mean_forest_pred_train[,sample_counter] <- forest_model_mean$get_cached_forest_predictions()
+                }
             }
             if (include_variance_forest) {
                 forest_model_variance$sample_one_iteration(
                     forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_variance, 
                     active_forest = active_forest_variance, rng = rng, forest_model_config = forest_model_config_variance, 
                     global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
                 )
+                
+                # Cache predictions
+                if (keep_sample) {
+                    variance_forest_pred_train[,sample_counter] <- forest_model_variance$get_cached_forest_predictions()
+                }
             }
             if (sample_sigma2_global) {
                 current_sigma2 <- sampleGlobalErrorVarianceOneIteration(outcome_train, forest_dataset_train, rng, a_global, b_global)
@@ -910,13 +922,21 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                         active_forest = active_forest_mean, rng = rng, forest_model_config = forest_model_config_mean, 
                         global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
                     )
+                    
+                    if (keep_sample) {
+                        mean_forest_pred_train[,sample_counter] <- forest_model_mean$get_cached_forest_predictions()
+                    }
                 }
                 if (include_variance_forest) {
                     forest_model_variance$sample_one_iteration(
                         forest_dataset = forest_dataset_train, residual = outcome_train, forest_samples = forest_samples_variance, 
                         active_forest = active_forest_variance, rng = rng, forest_model_config = forest_model_config_variance, 
                         global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
                     )
+                    
+                    if (keep_sample) {
+                        variance_forest_pred_train[,sample_counter] <- forest_model_variance$get_cached_forest_predictions()
+                    }
                 }
                 if (sample_sigma2_global) {
                     current_sigma2 <- sampleGlobalErrorVarianceOneIteration(outcome_train, forest_dataset_train, rng, a_global, b_global)
@@ -949,6 +969,12 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                 rfx_samples$delete_sample(0)
             }
         }
+        if (include_mean_forest) {
+            mean_forest_pred_train <- mean_forest_pred_train[,(num_gfr+1):ncol(mean_forest_pred_train)]
+        }
+        if (include_variance_forest) {
+            variance_forest_pred_train <- variance_forest_pred_train[,(num_gfr+1):ncol(variance_forest_pred_train)]
+        }
         if (sample_sigma2_global) {
             global_var_samples <- global_var_samples[(num_gfr+1):length(global_var_samples)]
         }
@@ -960,13 +986,15 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
 
     # Mean forest predictions
     if (include_mean_forest) {
-        y_hat_train <- forest_samples_mean$predict(forest_dataset_train)*y_std_train + y_bar_train
+        # y_hat_train <- forest_samples_mean$predict(forest_dataset_train)*y_std_train + y_bar_train
+        y_hat_train <- mean_forest_pred_train*y_std_train + y_bar_train
         if (has_test) y_hat_test <- forest_samples_mean$predict(forest_dataset_test)*y_std_train + y_bar_train
     }
 
     # Variance forest predictions
     if (include_variance_forest) {
-        sigma2_x_hat_train <- forest_samples_variance$predict(forest_dataset_train)
+        # sigma2_x_hat_train <- forest_samples_variance$predict(forest_dataset_train)
+        sigma2_x_hat_train <- variance_forest_pred_train
         if (has_test) sigma2_x_hat_test <- forest_samples_variance$predict(forest_dataset_test)
     }
 

R/model.R

@@ -126,6 +126,13 @@ ForestModel <- R6::R6Class(
             }
         }, 
 
+        #' @description
+        #' Extract an internally-cached prediction of a forest on the training dataset in a sampler.
+        #' @return Vector with as many elements as observations in the training dataset
+        get_cached_forest_predictions = function() {
+            get_cached_forest_predictions_cpp(self$tracker_ptr)
+        }, 
+        
         #' @description
         #' Propagates basis update through to the (full/partial) residual by iteratively 
         #' (a) adding back in the previous prediction of each tree, (b) recomputing predictions 

include/stochtree/partition_tracker.h

@@ -91,6 +91,10 @@ class ForestTracker {
   SampleNodeMapper* GetSampleNodeMapper() {return sample_node_mapper_.get();}
   UnsortedNodeSampleTracker* GetUnsortedNodeSampleTracker() {return unsorted_node_sample_tracker_.get();}
   SortedNodeSampleTracker* GetSortedNodeSampleTracker() {return sorted_node_sample_tracker_.get();}
+  int GetNumObservations() {return num_observations_;}
+  int GetNumTrees() {return num_trees_;}
+  int GetNumFeatures() {return num_features_;}
+  bool Initialized() {return initialized_;}
 
  private:
   /*! \brief Mapper from observations to predicted values summed over every tree in a forest */

src/py_stochtree.cpp

@@ -2147,6 +2147,7 @@ PYBIND11_MODULE(stochtree_cpp, m) {
     .def("ReconstituteTrackerFromForest", &ForestSamplerCpp::ReconstituteTrackerFromForest)
     .def("SampleOneIteration", &ForestSamplerCpp::SampleOneIteration)
     .def("InitializeForestModel", &ForestSamplerCpp::InitializeForestModel)
+    .def("GetCachedForestPredictions", &ForestSamplerCpp::GetCachedForestPredictions)
     .def("PropagateBasisUpdate", &ForestSamplerCpp::PropagateBasisUpdate)
     .def("PropagateResidualUpdate", &ForestSamplerCpp::PropagateResidualUpdate)
     .def("UpdateAlpha", &ForestSamplerCpp::UpdateAlpha)

src/sampler.cpp

@@ -284,6 +284,16 @@ cpp11::external_pointer<StochTree::ForestTracker> forest_tracker_cpp(cpp11::exte
     return cpp11::external_pointer<StochTree::ForestTracker>(tracker_ptr_.release());
 }
 
+[[cpp11::register]]
+cpp11::writable::doubles get_cached_forest_predictions_cpp(cpp11::external_pointer<StochTree::ForestTracker> tracker_ptr) {
+    int n_train = tracker_ptr->GetNumObservations();
+    cpp11::writable::doubles output(n_train);
+    for (int i = 0; i < n_train; i++) {
+        output[i] = tracker_ptr->GetSamplePrediction(i);
+    }
+    return output;
+}
+
 [[cpp11::register]]
 cpp11::writable::integers sample_without_replacement_integer_cpp(
     cpp11::integers population_vector,