Updated BART and BCF caching logic, added unit tests, and fixed bugs

Author: andrewherren

R/bart.R

@@ -760,7 +760,7 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                     global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
                 )
 
-                # Cache predictions
+                # Cache train set predictions since they are already computed during sampling
                 if (keep_sample) {
                     mean_forest_pred_train[,sample_counter] <- forest_model_mean$get_cached_forest_predictions()
                 }
@@ -772,7 +772,7 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                     global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
                 )
 
-                # Cache predictions
+                # Cache train set predictions since they are already computed during sampling
                 if (keep_sample) {
                     variance_forest_pred_train[,sample_counter] <- forest_model_variance$get_cached_forest_predictions()
                 }
@@ -923,6 +923,7 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                         global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
                     )
 
+                    # Cache train set predictions since they are already computed during sampling
                     if (keep_sample) {
                         mean_forest_pred_train[,sample_counter] <- forest_model_mean$get_cached_forest_predictions()
                     }
@@ -934,6 +935,7 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
                         global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
                     )
 
+                    # Cache train set predictions since they are already computed during sampling
                     if (keep_sample) {
                         variance_forest_pred_train[,sample_counter] <- forest_model_variance$get_cached_forest_predictions()
                     }

R/bcf.R

@@ -885,6 +885,8 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
     if (sample_sigma2_global) global_var_samples <- rep(NA, num_retained_samples)
     if (sample_sigma2_leaf_mu) leaf_scale_mu_samples <- rep(NA, num_retained_samples)
     if (sample_sigma2_leaf_tau) leaf_scale_tau_samples <- rep(NA, num_retained_samples)
+    muhat_train_raw <- matrix(NA_real_, nrow(X_train), num_retained_samples)
+    if (include_variance_forest) sigma2_x_train_raw <- matrix(NA_real_, nrow(X_train), num_retained_samples)
     sample_counter <- 0
 
     # Prepare adaptive coding structure
@@ -997,6 +999,11 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                 global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
             )
 
+            # Cache train set predictions since they are already computed during sampling
+            if (keep_sample) {
+                muhat_train_raw[,sample_counter] <- forest_model_mu$get_cached_forest_predictions()
+            }
+            
             # Sample variance parameters (if requested)
             if (sample_sigma2_global) {
                 current_sigma2 <- sampleGlobalErrorVarianceOneIteration(outcome_train, forest_dataset_train, rng, a_global, b_global)
@@ -1016,6 +1023,10 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                 global_model_config = global_model_config, keep_forest = keep_sample, gfr = TRUE
             )
 
+            # Cannot cache train set predictions for tau because the cached predictions in the 
+            # tracking data structures are pre-multiplied by the basis (treatment)
+            # ...
+
             # Sample coding parameters (if requested)
             if (adaptive_coding) {
                 # Estimate mu(X) and tau(X) and compute y - mu(X)
@@ -1060,6 +1071,11 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                     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 train set predictions since they are already computed during sampling
+                if (keep_sample) {
+                    sigma2_x_train_raw[,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)
@@ -1263,6 +1279,11 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                     global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
                 )
 
+                # Cache train set predictions since they are already computed during sampling
+                if (keep_sample) {
+                    muhat_train_raw[,sample_counter] <- forest_model_mu$get_cached_forest_predictions()
+                }
+                
                 # Sample variance parameters (if requested)
                 if (sample_sigma2_global) {
                     current_sigma2 <- sampleGlobalErrorVarianceOneIteration(outcome_train, forest_dataset_train, rng, a_global, b_global)
@@ -1282,6 +1303,10 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                     global_model_config = global_model_config, keep_forest = keep_sample, gfr = FALSE
                 )
 
+                # Cannot cache train set predictions for tau because the cached predictions in the 
+                # tracking data structures are pre-multiplied by the basis (treatment)
+                # ...
+                
                 # Sample coding parameters (if requested)
                 if (adaptive_coding) {
                     # Estimate mu(X) and tau(X) and compute y - mu(X)
@@ -1326,6 +1351,11 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
                         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
                     )
+                    
+                    # Cache train set predictions since they are already computed during sampling
+                    if (keep_sample) {
+                        sigma2_x_train_raw[,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)
@@ -1372,11 +1402,15 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
             b_1_samples <- b_1_samples[(num_gfr+1):length(b_1_samples)]
             b_0_samples <- b_0_samples[(num_gfr+1):length(b_0_samples)]
         }
+        muhat_train_raw <- muhat_train_raw[,(num_gfr+1):ncol(muhat_train_raw)]
+        if (include_variance_forest) {
+            sigma2_x_train_raw <- sigma2_x_train_raw[,(num_gfr+1):ncol(sigma2_x_train_raw)]
+        }
         num_retained_samples <- num_retained_samples - num_gfr
     }
 
     # Forest predictions
-    mu_hat_train <- forest_samples_mu$predict(forest_dataset_train)*y_std_train + y_bar_train
+    mu_hat_train <- muhat_train_raw*y_std_train + y_bar_train
     if (adaptive_coding) {
         tau_hat_train_raw <- forest_samples_tau$predict_raw(forest_dataset_train)
         tau_hat_train <- t(t(tau_hat_train_raw) * (b_1_samples - b_0_samples))*y_std_train
@@ -1395,7 +1429,7 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
         y_hat_test <- mu_hat_test + tau_hat_test * as.numeric(Z_test)
     }
     if (include_variance_forest) {
-        sigma2_x_hat_train <- forest_samples_variance$predict(forest_dataset_train)
+        sigma2_x_hat_train <- exp(sigma2_x_train_raw)
         if (has_test) sigma2_x_hat_test <- forest_samples_variance$predict(forest_dataset_test)
     }
 

src/container.cpp

@@ -206,6 +206,7 @@ void ForestContainer::append_from_json(const json& forest_container_json) {
   CHECK_EQ(this->num_trees_, forest_container_json.at("num_trees"));
   CHECK_EQ(this->output_dimension_, forest_container_json.at("output_dimension"));
   CHECK_EQ(this->is_leaf_constant_, forest_container_json.at("is_leaf_constant"));
+  CHECK_EQ(this->is_exponentiated_, forest_container_json.at("is_exponentiated"));
   CHECK_EQ(this->initialized_, forest_container_json.at("initialized"));
   int new_num_samples = forest_container_json.at("num_samples");
 
@@ -215,8 +216,7 @@ void ForestContainer::append_from_json(const json& forest_container_json) {
   for (int i = 0; i < forest_container_json.at("num_samples"); i++) {
     forest_ind = this->num_samples_ + i;
     forest_label = "forest_" + std::to_string(i);
-    // forests_[forest_ind] = std::make_unique<TreeEnsemble>(this->num_trees_, this->output_dimension_, this->is_leaf_constant_);
-    forests_.push_back(std::make_unique<TreeEnsemble>(this->num_trees_, this->output_dimension_, this->is_leaf_constant_));
+    forests_.push_back(std::make_unique<TreeEnsemble>(this->num_trees_, this->output_dimension_, this->is_leaf_constant_, this->is_exponentiated_));
     forests_[forest_ind]->from_json(forest_container_json.at(forest_label));
   }
   this->num_samples_ += new_num_samples;

stochtree/bart.py

@@ -1191,6 +1191,7 @@ def sample(
                         True,
                     )
 
+                    # Cache train set predictions since they are already computed during sampling
                     if keep_sample:
                         yhat_train_raw[:,sample_counter] = forest_sampler_mean.get_cached_forest_predictions()
 
@@ -1208,6 +1209,7 @@ def sample(
                         True,
                     )
 
+                    # Cache train set predictions since they are already computed during sampling
                     if keep_sample:
                         sigma2_x_train_raw[:,sample_counter] = forest_sampler_variance.get_cached_forest_predictions()
 
@@ -1389,6 +1391,9 @@ def sample(
                             False,
                         )
 
+                        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(
@@ -1403,6 +1408,9 @@ def sample(
                             False,
                         )
 
+                        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(
@@ -1493,16 +1501,15 @@ def sample(
                     self.y_hat_test = rfx_preds_test
 
         if self.include_variance_forest:
-            sigma2_x_train_raw = np.exp(sigma2_x_train_raw)
             if self.sample_sigma2_global:
-                self.sigma2_x_train = sigma2_x_train_raw
+                self.sigma2_x_train = np.empty_like(sigma2_x_train_raw)
                 for i in range(self.num_samples):
                     self.sigma2_x_train[:, i] = (
-                        sigma2_x_train_raw[:, i] * self.global_var_samples[i]
+                        np.exp(sigma2_x_train_raw[:, i]) * self.global_var_samples[i]
                     )
             else:
                 self.sigma2_x_train = (
-                    sigma2_x_train_raw * self.sigma2_init * self.y_std * self.y_std
+                    np.exp(sigma2_x_train_raw) * self.sigma2_init * self.y_std * self.y_std
                 )
             if self.has_test:
                 sigma2_x_test_raw = (
@@ -1628,14 +1635,14 @@ def predict(
                 )
             )
             if self.sample_sigma2_global:
-                variance_pred = variance_pred_raw
+                variance_pred = np.empty_like(variance_pred_raw)
                 for i in range(self.num_samples):
-                    variance_pred[:, i] = np.sqrt(
+                    variance_pred[:, i] = (
                         variance_pred_raw[:, i] * self.global_var_samples[i]
                     )
             else:
                 variance_pred = (
-                    np.sqrt(variance_pred_raw * self.sigma2_init) * self.y_std
+                    variance_pred_raw * self.sigma2_init * self.y_std * self.y_std
                 )
 
         has_mean_predictions = self.include_mean_forest or self.has_rfx
@@ -1817,7 +1824,7 @@ def predict_variance(self, covariates: np.array) -> np.array:
             pred_dataset.dataset_cpp
         )
         if self.sample_sigma2_global:
-            variance_pred = variance_pred_raw
+            variance_pred = np.empty_like(variance_pred_raw)
             for i in range(self.num_samples):
                 variance_pred[:, i] = (
                     variance_pred_raw[:, i] * self.global_var_samples[i]
@@ -2024,11 +2031,11 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
             for i in range(len(json_object_list)):
                 if i == 0:
                     self.forest_container_variance.forest_container_cpp.LoadFromJson(
-                        json_object_list[i].json_cpp, "forest_1"
+                        json_object_list[i].json_cpp, "forest_0"
                     )
                 else:
                     self.forest_container_variance.forest_container_cpp.AppendFromJson(
-                        json_object_list[i].json_cpp, "forest_1"
+                        json_object_list[i].json_cpp, "forest_0"
                     )
 
         # Unpack random effects
@@ -2053,13 +2060,19 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
         self.num_gfr = json_object_default.get_integer("num_gfr")
         self.num_burnin = json_object_default.get_integer("num_burnin")
         self.num_mcmc = json_object_default.get_integer("num_mcmc")
-        self.num_samples = json_object_default.get_integer("num_samples")
         self.num_basis = json_object_default.get_integer("num_basis")
         self.has_basis = json_object_default.get_boolean("requires_basis")
         self.probit_outcome_model = json_object_default.get_boolean(
             "probit_outcome_model"
         )
 
+        # Unpack number of samples
+        for i in range(len(json_object_list)):
+            if i == 0:
+                self.num_samples = json_object_list[i].get_integer("num_samples")
+            else:
+                self.num_samples += json_object_list[i].get_integer("num_samples")
+
         # Unpack parameter samples
         if self.sample_sigma2_global:
             for i in range(len(json_object_list)):