Merge pull request #233 from StochasticTree/predict-python-hotfix

Hotfix for BCF predictions when only tau(X) is requested

Author: andrewherren

R/bcf.R

@@ -3061,8 +3061,10 @@ predict.bcfmodel <- function(
         t(tau_hat_raw) * (object$b_1_samples - object$b_0_samples)
       ) *
         y_std
-      control_adj <- t(t(tau_hat_raw) * object$b_0_samples) * y_std
-      mu_hat_forest <- mu_hat_forest + control_adj
+      if (predict_mu_forest || predict_mu_forest_intermediate) {
+        control_adj <- t(t(tau_hat_raw) * object$b_0_samples) * y_std
+        mu_hat_forest <- mu_hat_forest + control_adj
+      }
     } else {
       tau_hat_forest <- object$forests_tau$predict_raw(forest_dataset_pred) *
         y_std

stochtree/bcf.py

@@ -2604,12 +2604,13 @@ def predict(
                 adaptive_coding_weights = np.expand_dims(
                     self.b1_samples - self.b0_samples, axis=(0, 2)
                 )
-                b0_weights = np.expand_dims(
-                    self.b0_samples, axis=(0, 2)
-                )
-                control_adj = tau_raw * b0_weights * self.y_std
+                if predict_mu_forest or predict_mu_forest_intermediate:
+                    b0_weights = np.expand_dims(
+                        self.b0_samples, axis=(0, 2)
+                    )
+                    control_adj = tau_raw * b0_weights * self.y_std
+                    mu_x_forest = mu_x_forest + np.squeeze(control_adj)
                 tau_raw = tau_raw * adaptive_coding_weights
-                mu_x_forest = mu_x_forest + np.squeeze(control_adj)
             tau_x_forest = np.squeeze(tau_raw * self.y_std)
             if Z.shape[1] > 1:
                 treatment_term = np.multiply(