Updated ensemble predict method to skip over a dataset basis for prediction if not used in a model

Author: andrewherren

R/bcf.R

@@ -770,8 +770,8 @@ bcf <- function(X_train, Z_train, y_train, pi_train = NULL, group_ids_train = NU
         "outcome_mean" = y_bar_train,
         "outcome_scale" = y_std_train, 
         "num_covariates" = ncol(X_train),
-        "num_prognostic_covariates" = ncol(X_train_mu),
-        "num_treatment_covariates" = ncol(X_train_tau),
+        "num_prognostic_covariates" = sum(variable_weights_mu > 0),
+        "num_treatment_covariates" = sum(variable_weights_tau > 0),
         "treatment_dim" = ncol(Z_train), 
         "propensity_covariate" = propensity_covariate, 
         "binary_treatment" = binary_treatment, 

include/stochtree/ensemble.h

@@ -81,12 +81,11 @@ class TreeEnsemble {
 
   inline void PredictInplace(ForestDataset& dataset, std::vector<double> &output, 
                              int tree_begin, int tree_end, data_size_t offset = 0) {
-    if (dataset.HasBasis()) {
-      CHECK(!is_leaf_constant_);
-      PredictInplace(dataset.GetCovariates(), dataset.GetBasis(), output, tree_begin, tree_end, offset);
-    } else {
-      CHECK(is_leaf_constant_);
+    if (is_leaf_constant_) {
       PredictInplace(dataset.GetCovariates(), output, tree_begin, tree_end, offset);
+    } else {
+      CHECK(dataset.HasBasis());
+      PredictInplace(dataset.GetCovariates(), dataset.GetBasis(), output, tree_begin, tree_end, offset);
     }
   }