Updated python predict methods and tests

Author: andrewherren

stochtree/bart.py

@@ -1803,13 +1803,6 @@ def predict(
             #     rfx_predictions = np.mean(rfx_predictions, axis = 1)
 
         # Combine into y hat predictions
-        if predict_y_hat and has_mean_forest and has_rfx:
-            y_hat = mean_forest_predictions + rfx_predictions
-        elif predict_y_hat and has_mean_forest:
-            y_hat = mean_forest_predictions
-        elif predict_y_hat and has_rfx:
-            y_hat = rfx_predictions
-
         if probability_scale:
             if predict_y_hat and has_mean_forest and has_rfx:
                 y_hat = norm.ppf(mean_forest_predictions + rfx_predictions)

stochtree/bcf.py

@@ -2400,14 +2400,28 @@ def predict(
         forest_dataset_test.add_covariates(X_combined)
         forest_dataset_test.add_basis(Z)
 
-        # Compute predicted outcome and decomposed outcome model terms
+        # Compute predictions from the variance forest (if included)
+        if predict_variance_forest:
+            sigma2_x_raw = self.forest_container_variance.forest_container_cpp.Predict(
+                forest_dataset_test.dataset_cpp
+            )
+            if self.sample_sigma2_global:
+                sigma2_x = np.empty_like(sigma2_x_raw)
+                for i in range(self.num_samples):
+                    sigma2_x[:, i] = sigma2_x_raw[:, i] * self.global_var_samples[i]
+            else:
+                sigma2_x = sigma2_x_raw * self.sigma2_init * self.y_std * self.y_std
+            if predict_mean:
+                sigma2_x = np.mean(sigma2_x, axis=1)
+
+        # Prognostic forest predictions
         if predict_mu_forest or predict_mu_forest_intermediate:
             mu_raw = self.forest_container_mu.forest_container_cpp.Predict(
                 forest_dataset_test.dataset_cpp
             )
             mu_x = mu_raw * self.y_std + self.y_bar
-            if predict_mean:
-                mu_x = np.mean(mu_x, axis=1)
+        
+        # Treatment effect forest predictions
         if predict_tau_forest or predict_tau_forest_intermediate:
             tau_raw = self.forest_container_tau.forest_container_cpp.PredictRaw(
                 forest_dataset_test.dataset_cpp
@@ -2422,43 +2436,65 @@ def predict(
                 treatment_term = np.multiply(
                     np.atleast_3d(Z).swapaxes(1, 2), tau_x
                 ).sum(axis=2)
-                if predict_mean:
-                    treatment_term = np.mean(treatment_term, axis=1)
-                    tau_x = np.mean(tau_x, axis=2)
             else:
                 treatment_term = Z * np.squeeze(tau_x)
-                if predict_mean:
-                    treatment_term = np.mean(treatment_term, axis=1)
-                    tau_x = np.mean(tau_x, axis=1)
 
+        # Random effects predictions
         if predict_rfx or predict_rfx_intermediate:
             rfx_preds = (
                 self.rfx_container.predict(rfx_group_ids, rfx_basis) * self.y_std
             )
             if predict_mean:
                 rfx_preds = np.mean(rfx_preds, axis=1)
 
+        # Combine into y hat predictions
         if predict_y_hat and has_mu_forest and has_rfx:
             y_hat = mu_x + treatment_term + rfx_preds
         elif predict_y_hat and has_mu_forest:
             y_hat = mu_x + treatment_term
         elif predict_y_hat and has_rfx:
             y_hat = rfx_preds
-
-        # Compute predictions from the variance forest (if included)
-        if predict_variance_forest:
-            sigma2_x_raw = self.forest_container_variance.forest_container_cpp.Predict(
-                forest_dataset_test.dataset_cpp
-            )
-            if self.sample_sigma2_global:
-                sigma2_x = np.empty_like(sigma2_x_raw)
-                for i in range(self.num_samples):
-                    sigma2_x[:, i] = sigma2_x_raw[:, i] * self.global_var_samples[i]
+        
+        needs_mean_term_preds = predict_y_hat or \
+            predict_mu_forest or \
+            predict_tau_forest or \
+            predict_rfx
+        if needs_mean_term_preds:
+            if probability_scale:
+                if has_rfx:
+                    if predict_y_hat:
+                        y_hat = norm.cdf(mu_x + treatment_term + rfx_preds)
+                    if predict_rfx:
+                        rfx_preds = norm.cdf(rfx_preds)
+                else:
+                    if predict_y_hat:
+                        y_hat = norm.cdf(mu_x + treatment_term)
+                if predict_mu_forest:
+                    mu_x = norm.cdf(mu_x)
+                if predict_tau_forest:
+                    tau_x = norm.cdf(tau_x)
             else:
-                sigma2_x = sigma2_x_raw * self.sigma2_init * self.y_std * self.y_std
-            if predict_mean:
-                sigma2_x = np.mean(sigma2_x, axis=1)
+                if has_rfx:
+                    if predict_y_hat:
+                        y_hat = mu_x + treatment_term + rfx_preds
+                else:
+                    if predict_y_hat:
+                        y_hat = mu_x + treatment_term
 
+        # Collapse to posterior mean predictions if requested
+        if predict_mean:
+            if predict_mu_forest:
+                mu_x = np.mean(mu_x, axis=1)
+            if predict_tau_forest:
+                if Z.shape[1] > 1:
+                    tau_x = np.mean(tau_x, axis=2)
+                else:
+                    tau_x = np.mean(tau_x, axis=1)
+            if predict_rfx:
+                rfx_preds = np.mean(rfx_preds, axis=1)
+            if predict_y_hat:
+                y_hat = np.mean(y_hat, axis=1)
+        
         if predict_count == 1:
             if predict_y_hat:
                 return y_hat
@@ -2754,6 +2790,9 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
         self.internal_propensity_model = json_object_default.get_boolean(
             "internal_propensity_model"
         )
+        self.probit_outcome_model = json_object_default.get_boolean(
+            "probit_outcome_model"
+        )
 
         # Unpack number of samples
         for i in range(len(json_object_list)):

test/python/test_bcf.py

@@ -76,8 +76,8 @@ def test_binary_bcf(self):
         assert mu_hat.shape == (n_test, num_mcmc)
         assert y_hat.shape == (n_test, num_mcmc)
 
-        # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        # Check that we can predict just treatment effects
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF without test set and with propensity score
@@ -106,7 +106,7 @@ def test_binary_bcf(self):
         assert mu_hat.shape == (n_test, num_mcmc)
         assert y_hat.shape == (n_test, num_mcmc)
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF with test set and without propensity score
@@ -142,7 +142,7 @@ def test_binary_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF without test set and without propensity score
@@ -172,7 +172,7 @@ def test_binary_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, terms = "cate")
 
     def test_continuous_univariate_bcf(self):
         # RNG
@@ -245,7 +245,7 @@ def test_continuous_univariate_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run second BCF model with test set and propensity score
@@ -281,7 +281,7 @@ def test_continuous_univariate_bcf(self):
         assert y_hat_2.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat_2 = bcf_model_2.predict_tau(X_test, Z_test, pi_test)
+        tau_hat_2 = bcf_model_2.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat_2.shape == (n_test, num_mcmc)
 
         # Combine into a single model
@@ -336,7 +336,7 @@ def test_continuous_univariate_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF with test set and without propensity score
@@ -372,7 +372,7 @@ def test_continuous_univariate_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF without test set and without propensity score
@@ -402,7 +402,7 @@ def test_continuous_univariate_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, terms = "cate")
 
         # Run second BCF model with test set and propensity score
         bcf_model_2 = BCFModel()
@@ -430,7 +430,7 @@ def test_continuous_univariate_bcf(self):
         assert y_hat_2.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat_2 = bcf_model_2.predict_tau(X_test, Z_test)
+        tau_hat_2 = bcf_model_2.predict(X = X_test, Z = Z_test, terms = "cate")
         assert tau_hat_2.shape == (n_test, num_mcmc)
 
         # Combine into a single model
@@ -528,7 +528,7 @@ def test_multivariate_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc, treatment_dim)
 
         # Run BCF without test set and with propensity score
@@ -558,7 +558,7 @@ def test_multivariate_bcf(self):
         assert y_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc, treatment_dim)
 
         # Run BCF with test set and without propensity score
@@ -665,7 +665,7 @@ def test_binary_bcf_heteroskedastic(self):
         assert sigma2_x_hat.shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF without test set and with propensity score
@@ -715,7 +715,7 @@ def test_binary_bcf_heteroskedastic(self):
         )
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test, pi_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, propensity = pi_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF with test set and without propensity score
@@ -752,7 +752,7 @@ def test_binary_bcf_heteroskedastic(self):
         assert bcf_preds['variance_forest_predictions'].shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, terms = "cate")
         assert tau_hat.shape == (n_test, num_mcmc)
 
         # Run BCF without test set and without propensity score
@@ -781,7 +781,7 @@ def test_binary_bcf_heteroskedastic(self):
         assert bcf_preds['y_hat'].shape == (n_test, num_mcmc)
 
         # Check treatment effect prediction method
-        tau_hat = bcf_model.predict_tau(X_test, Z_test)
+        tau_hat = bcf_model.predict(X = X_test, Z = Z_test, terms = "cate")
 
     def test_bcf_rfx_parameters(self):
         # RNG

test/python/test_predict.py

@@ -263,10 +263,6 @@ def test_bart_prediction(self):
     def test_bcf_prediction(self):
         # Generate data and test/train split
         rng = np.random.default_rng(1234)
-        
-
-        # Convert the R code down below to Python
-        rng = np.random.default_rng(1234)
         n = 100
         g = lambda x: np.where(x[:, 4] == 1, 2, np.where(x[:, 4] == 2, -1, -4))
         x1 = rng.normal(size=n)
@@ -328,7 +324,7 @@ def g(x5):
             num_mcmc = 10
         )
 
-        # Check that the default predict method returns a list
+        # Check that the default predict method returns a dictionary
         pred = bcf_model.predict(X=X_test, Z=Z_test, propensity=pi_x_test)
         y_hat_posterior_test = pred['y_hat']
         assert y_hat_posterior_test.shape == (20, 10)