Updated python BCF predictions

Author: andrewherren

stochtree/bcf.py

@@ -2267,7 +2267,12 @@ def sample(
             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_train = self.tau_hat_train * b0_weights * self.y_std
             self.tau_hat_train = self.tau_hat_train * adaptive_coding_weights
+            self.mu_hat_train = self.mu_hat_train + np.squeeze(control_adj_train)
         self.tau_hat_train = np.squeeze(self.tau_hat_train * self.y_std)
         if self.multivariate_treatment:
             treatment_term_train = np.multiply(
@@ -2289,7 +2294,12 @@ def sample(
                 adaptive_coding_weights_test = 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_test = self.tau_hat_test * b0_weights * self.y_std
                 self.tau_hat_test = self.tau_hat_test * adaptive_coding_weights_test
+                self.mu_hat_test = self.mu_hat_test + np.squeeze(control_adj_test)
             self.tau_hat_test = np.squeeze(self.tau_hat_test * self.y_std)
             if self.multivariate_treatment:
                 treatment_term_test = np.multiply(
@@ -2594,7 +2604,12 @@ 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
                 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(