Update Python and R train set prediction caching logic

Author: andrewherren

R/bart.R

@@ -994,7 +994,7 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
     # Variance forest predictions
     if (include_variance_forest) {
         # sigma2_x_hat_train <- forest_samples_variance$predict(forest_dataset_train)
-        sigma2_x_hat_train <- variance_forest_pred_train
+        sigma2_x_hat_train <- exp(variance_forest_pred_train)
         if (has_test) sigma2_x_hat_test <- forest_samples_variance$predict(forest_dataset_test)
     }
 

src/py_stochtree.cpp

@@ -1166,6 +1166,16 @@ class ForestSamplerCpp {
     }
   }
 
+  py::array_t<double> GetCachedForestPredictions() {
+    int n_train = tracker_->GetNumObservations();
+    auto output = py::array_t<double>(py::detail::any_container<py::ssize_t>({n_train}));
+    auto accessor = output.mutable_unchecked<1>();
+    for (size_t i = 0; i < n_train; i++) {
+      accessor(i) = tracker_->GetSamplePrediction(i);
+    }
+    return output;
+  }
+
   void PropagateBasisUpdate(ForestDatasetCpp& dataset, ResidualCpp& residual, ForestCpp& forest) {
     // Perform the update operation
     StochTree::UpdateResidualNewBasis(*tracker_, *(dataset.GetDataset()), *(residual.GetData()), forest.GetEnsemble());

stochtree/bart.py

@@ -1005,6 +1005,10 @@ def sample(
             self.global_var_samples = np.empty(self.num_samples, dtype=np.float64)
         if sample_sigma2_leaf:
             self.leaf_scale_samples = np.empty(self.num_samples, dtype=np.float64)
+        if self.include_mean_forest:
+            yhat_train_raw = np.empty((self.n_train, self.num_samples), dtype=np.float64)
+        if self.include_variance_forest:
+            sigma2_x_train_raw = np.empty((self.n_train, self.num_samples), dtype=np.float64)
         sample_counter = -1
 
         # Forest Dataset (covariates and optional basis)
@@ -1187,6 +1191,9 @@ def sample(
                         True,
                     )
 
+                    if keep_sample:
+                        yhat_train_raw[:,sample_counter] = forest_sampler_mean.get_cached_forest_predictions()
+
                 # Sample the variance forest
                 if self.include_variance_forest:
                     forest_sampler_variance.sample_one_iteration(
@@ -1201,6 +1208,9 @@ def sample(
                         True,
                     )
 
+                    if keep_sample:
+                        sigma2_x_train_raw[:,sample_counter] = forest_sampler_variance.get_cached_forest_predictions()
+
                 # Sample variance parameters (if requested)
                 if self.sample_sigma2_global:
                     current_sigma2 = global_var_model.sample_one_iteration(
@@ -1441,6 +1451,10 @@ def sample(
                 self.global_var_samples = self.global_var_samples[num_gfr:]
             if self.sample_sigma2_leaf:
                 self.leaf_scale_samples = self.leaf_scale_samples[num_gfr:]
+            if self.include_mean_forest:
+                yhat_train_raw = yhat_train_raw[:,num_gfr:]
+            if self.include_variance_forest:
+                sigma2_x_train_raw = sigma2_x_train_raw[:,num_gfr:]
             self.num_samples -= num_gfr
 
         # Store predictions
@@ -1451,9 +1465,6 @@ def sample(
             self.leaf_scale_samples = self.leaf_scale_samples
 
         if self.include_mean_forest:
-            yhat_train_raw = self.forest_container_mean.forest_container_cpp.Predict(
-                forest_dataset_train.dataset_cpp
-            )
             self.y_hat_train = yhat_train_raw * self.y_std + self.y_bar
             if self.has_test:
                 yhat_test_raw = self.forest_container_mean.forest_container_cpp.Predict(
@@ -1482,11 +1493,7 @@ def sample(
                     self.y_hat_test = rfx_preds_test
 
         if self.include_variance_forest:
-            sigma2_x_train_raw = (
-                self.forest_container_variance.forest_container_cpp.Predict(
-                    forest_dataset_train.dataset_cpp
-                )
-            )
+            sigma2_x_train_raw = np.exp(sigma2_x_train_raw)
             if self.sample_sigma2_global:
                 self.sigma2_x_train = sigma2_x_train_raw
                 for i in range(self.num_samples):

stochtree/sampler.py

@@ -266,6 +266,17 @@ def propagate_basis_update(
         self.forest_sampler_cpp.PropagateBasisUpdate(
             dataset.dataset_cpp, residual.residual_cpp, forest.forest_cpp
         )
+    
+    def get_cached_forest_predictions(self) -> np.array:
+        """
+        Extract an internally-cached prediction of a forest on the training dataset in a sampler.
+
+        Returns
+        ----------
+        np.array
+            Numpy 1D array with as many elements as observations in the training dataset
+        """
+        return self.forest_sampler_cpp.GetCachedForestPredictions()
 
     def update_alpha(self, alpha: float) -> None:
         """