Refactored predict method for BART and BCF to include labels for prediction components

andrewherren · andrewherren · commit 73bd6350392b · 2025-08-11T18:21:52.000-05:00
diff --git a/R/bart.R b/R/bart.R
@@ -1276,15 +1276,23 @@ predict.bartmodel <- function(object, X, leaf_basis = NULL, rfx_group_ids = NULL
     result <- list()
     if ((object$model_params$has_rfx) || (object$model_params$include_mean_forest)) {
         result[["y_hat"]] = y_hat
+    } else {
+        result[["y_hat"]] <- NULL
     }
     if (object$model_params$include_mean_forest) {
         result[["mean_forest_predictions"]] = mean_forest_predictions
+    } else {
+        result[["mean_forest_predictions"]] <- NULL
     }
     if (object$model_params$has_rfx) {
         result[["rfx_predictions"]] = rfx_predictions
+    } else {
+        result[["rfx_predictions"]] <- NULL
     }
     if (object$model_params$include_variance_forest) {
         result[["variance_forest_predictions"]] = variance_forest_predictions
+    } else {
+        result[["variance_forest_predictions"]] <- NULL
     }
     return(result)
 }
diff --git a/R/bcf.R b/R/bcf.R
@@ -1800,10 +1800,14 @@ predict.bcfmodel <- function(object, X, Z, propensity = NULL, rfx_group_ids = NU
         "y_hat" = y_hat
     )
     if (object$model_params$has_rfx) {
-        result[["rfx_predictions"]] = rfx_predictions
+        result[["rfx_predictions"]] <- rfx_predictions
+    } else {
+        result[["rfx_predictions"]] <- NULL
     }
     if (object$model_params$include_variance_forest) {
-        result[["variance_forest_predictions"]] = variance_forest_predictions
+        result[["variance_forest_predictions"]] <- variance_forest_predictions
+    } else {
+        result[["variance_forest_predictions"]] <- NULL
     }
     return(result)
 }
diff --git a/demo/debug/multi_chain.py b/demo/debug/multi_chain.py
@@ -118,11 +118,15 @@ def outcome_mean(X, W):
 )
 
 # Inspect the model outputs
-y_hat_mcmc_2 = bart_model_2.predict(X_test, basis_test)
+bart_preds_2 = bart_model_2.predict(X_test, basis_test)
+y_hat_mcmc_2 = bart_preds_2['y_hat']
 y_avg_mcmc_2 = np.squeeze(y_hat_mcmc_2).mean(axis=1, keepdims=True)
-y_hat_mcmc_3 = bart_model_3.predict(X_test, basis_test)
+y_avg_mcmc_2 = np.squeeze(y_hat_mcmc_2).mean(axis=1, keepdims=True)
+bart_preds_3 = bart_model_3.predict(X_test, basis_test)
+y_hat_mcmc_3 = bart_preds_3['y_hat']
 y_avg_mcmc_3 = np.squeeze(y_hat_mcmc_3).mean(axis=1, keepdims=True)
-y_hat_mcmc_4 = bart_model_4.predict(X_test, basis_test)
+bart_preds_4 = bart_model_4.predict(X_test, basis_test)
+y_hat_mcmc_4 = bart_preds_4['y_hat']
 y_avg_mcmc_4 = np.squeeze(y_hat_mcmc_4).mean(axis=1, keepdims=True)
 y_df = pd.DataFrame(
     np.concatenate(
diff --git a/demo/debug/parallel_multi_chain.py b/demo/debug/parallel_multi_chain.py
@@ -145,7 +145,8 @@ def outcome_mean(X, W):
     )
 
     # Inspect the model outputs
-    y_hat_mcmc = combined_bart.predict(X_test, basis_test)
+    bart_preds = combined_bart.predict(X_test, basis_test)
+    y_hat_mcmc = bart_preds['y_hat']
     y_avg_mcmc = np.squeeze(y_hat_mcmc).mean(axis=1, keepdims=True)
     y_df = pd.DataFrame(
         np.concatenate((y_avg_mcmc, np.expand_dims(y_test, axis=1)), axis=1),
diff --git a/demo/debug/random_effects.py b/demo/debug/random_effects.py
@@ -75,7 +75,8 @@ def outcome_mean(group_labels, basis):
 rfx_model.sample(rfx_dataset, outcome, rfx_tracker, rfx_container, True, 1.0, cpp_rng)
 
 # Inspect the samples
-rfx_preds = rfx_container.predict(group_labels, basis) * y_std + y_bar
+bart_preds = rfx_container.predict(group_labels, basis)
+rfx_preds = bart_preds['y_hat'] * y_std + y_bar
 rfx_comparison_df = pd.DataFrame(
     np.concatenate((rfx_preds, np.expand_dims(rfx_term, axis=1)), axis=1),
     columns=["Predicted", "Actual"],
diff --git a/demo/debug/rfx_serialization.py b/demo/debug/rfx_serialization.py
@@ -60,11 +60,13 @@ def rfx_mean(group_labels, basis):
                   rfx_basis_train=basis, num_gfr=10, num_mcmc=10)
 
 # Extract predictions from the sampler
-y_hat_orig = bart_orig.predict(X, W, group_labels, basis)
+bart_preds_orig = bart_orig.predict(X, W, group_labels, basis)
+y_hat_orig = bart_preds_orig['y_hat']
 
 # "Round-trip" the model to JSON string and back and check that the predictions agree
 bart_json_string = bart_orig.to_json()
 bart_reloaded = BARTModel()
 bart_reloaded.from_json(bart_json_string)
-y_hat_reloaded = bart_reloaded.predict(X, W, group_labels, basis)
+bart_preds_reloaded = bart_reloaded.predict(X, W, group_labels, basis)
+y_hat_reloaded = bart_preds_reloaded['y_hat']
 np.testing.assert_almost_equal(y_hat_orig, y_hat_reloaded)
diff --git a/demo/debug/serialization.py b/demo/debug/serialization.py
@@ -98,11 +98,13 @@ def outcome_mean(X, W):
     global_var_samples[i+1] = global_var_model.sample_one_iteration(residual, cpp_rng, a_global, b_global)
 
 # Extract predictions from the sampler
-y_hat_orig = forest_container.predict(dataset)
+bart_preds_orig = forest_container.predict(dataset)
+y_hat_orig = bart_preds_orig['y_hat']
 
 # "Round-trip" the forest to JSON string and back and check that the predictions agree
 forest_json_string = forest_container.dump_json_string()
 forest_container_reloaded = ForestContainer(num_trees, W.shape[1], False, False)
 forest_container_reloaded.load_from_json_string(forest_json_string)
-y_hat_reloaded = forest_container_reloaded.predict(dataset)
+bart_preds_reloaded = forest_container_reloaded.predict(dataset)
+y_hat_reloaded = bart_preds_reloaded['y_hat']
 np.testing.assert_almost_equal(y_hat_orig, y_hat_reloaded)
diff --git a/demo/notebooks/prototype_interface.ipynb b/demo/notebooks/prototype_interface.ipynb
@@ -341,7 +341,8 @@
    "outputs": [],
    "source": [
     "# Forest predictions\n",
-    "forest_preds = forest_container.predict(dataset) * y_std + y_bar\n",
+    "bart_preds = forest_container.predict(dataset)\n",
+    "forest_preds = bart_preds['y_hat'] * y_std + y_bar\n",
     "forest_preds_gfr = forest_preds[:, :num_warmstart]\n",
     "forest_preds_mcmc = forest_preds[:, num_warmstart:num_samples]\n",
     "\n",
@@ -1101,7 +1102,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "venv",
+   "display_name": "venv (3.12.9)",
    "language": "python",
    "name": "python3"
   },
diff --git a/demo/notebooks/serialization.ipynb b/demo/notebooks/serialization.ipynb
@@ -241,7 +241,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "y_hat_deserialized = bart_model_deserialized.predict(X_test, basis_test)\n",
+    "bart_preds_deserialized = bart_model_deserialized.predict(X_test, basis_test)\n",
+    "y_hat_deserialized = bart_preds_deserialized['y_hat']\n",
     "y_avg_mcmc_deserialized = np.squeeze(y_hat_deserialized).mean(axis=1, keepdims=True)\n",
     "y_df = pd.DataFrame(\n",
     "    np.concatenate((y_avg_mcmc, y_avg_mcmc_deserialized), axis=1),\n",
@@ -325,7 +326,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "y_hat_file_deserialized = bart_model_file_deserialized.predict(X_test, basis_test)\n",
+    "bart_preds_file_deserialized = bart_model_file_deserialized.predict(X_test, basis_test)\n",
+    "y_hat_file_deserialized = bart_preds_file_deserialized['y_hat']\n",
     "y_avg_mcmc_file_deserialized = np.squeeze(y_hat_file_deserialized).mean(\n",
     "    axis=1, keepdims=True\n",
     ")\n",
@@ -381,7 +383,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "venv",
+   "display_name": "venv (3.12.9)",
    "language": "python",
    "name": "python3"
   },
diff --git a/stochtree/bart.py b/stochtree/bart.py
@@ -1691,11 +1691,11 @@ def predict(
 
         has_mean_predictions = self.include_mean_forest or self.has_rfx
         if has_mean_predictions and self.include_variance_forest:
-            return (mean_pred, variance_pred)
+            return {"y_hat": mean_pred, "variance_forest_predictions": variance_pred}
         elif has_mean_predictions and not self.include_variance_forest:
-            return mean_pred
+            return {"y_hat": mean_pred, "variance_forest_predictions": None}
         elif not has_mean_predictions and self.include_variance_forest:
-            return variance_pred
+            return {"y_hat": None, "variance_forest_predictions": variance_pred}
 
     def predict_mean(
         self,
diff --git a/stochtree/bcf.py b/stochtree/bcf.py
@@ -2239,23 +2239,6 @@ def predict_tau(
         # Data checks
         if Z.shape[0] != X.shape[0]:
             raise ValueError("X and Z must have the same number of rows")
-        if propensity is not None:
-            if propensity.shape[0] != X.shape[0]:
-                raise ValueError("X and propensity must have the same number of rows")
-        else:
-            if self.propensity_covariate == "tau":
-                if not self.internal_propensity_model:
-                    raise ValueError(
-                        "Propensity scores not provided, but no propensity model was trained during sampling"
-                    )
-                else:
-                    propensity = np.mean(
-                        self.bart_propensity_model.predict(X), axis=1, keepdims=True
-                    )
-            else:
-                # Dummy propensities if not provided but also not needed
-                propensity = np.ones(X.shape[0])
-                propensity = np.expand_dims(propensity, 1)
 
         # Covariate preprocessing
         if not self._covariate_preprocessor._check_is_fitted():
@@ -2277,6 +2260,22 @@ def predict_tau(
                 covariates_processed = X
         else:
             covariates_processed = self._covariate_preprocessor.transform(X)
+        
+        # Handle propensities
+        if propensity is not None:
+            if propensity.shape[0] != X.shape[0]:
+                raise ValueError("X and propensity must have the same number of rows")
+        else:
+            if self.propensity_covariate != "none":
+                if not self.internal_propensity_model:
+                    raise ValueError(
+                        "Propensity scores not provided, but no propensity model was trained during sampling"
+                    )
+                else:
+                    internal_propensity_preds = self.bart_propensity_model.predict(covariates_processed)
+                    propensity = np.mean(
+                        internal_propensity_preds['y_hat'], axis=1, keepdims=True
+                    )
 
         # Update covariates to include propensities if requested
         if self.propensity_covariate == "none":
@@ -2362,6 +2361,22 @@ def predict_variance(
                 covariates_processed = covariates
         else:
             covariates_processed = self._covariate_preprocessor.transform(covariates)
+        
+        # Handle propensities
+        if propensity is not None:
+            if propensity.shape[0] != covariates.shape[0]:
+                raise ValueError("X and propensity must have the same number of rows")
+        else:
+            if self.propensity_covariate != "none":
+                if not self.internal_propensity_model:
+                    raise ValueError(
+                        "Propensity scores not provided, but no propensity model was trained during sampling"
+                    )
+                else:
+                    internal_propensity_preds = self.bart_propensity_model.predict(covariates_processed)
+                    propensity = np.mean(
+                        internal_propensity_preds['y_hat'], axis=1, keepdims=True
+                    )
 
         # Update covariates to include propensities if requested
         if self.propensity_covariate == "none":
@@ -2402,7 +2417,7 @@ def predict(
         propensity: np.array = None,
         rfx_group_ids: np.array = None,
         rfx_basis: np.array = None,
-    ) -> tuple:
+    ) -> dict:
         """Predict outcome model components (CATE function and prognostic function) as well as overall outcome for every provided observation.
         Predicted outcomes are computed as `yhat = mu_x + Z*tau_x` where mu_x is a sample of the prognostic function and tau_x is a sample of the treatment effect (CATE) function.
 
@@ -2455,19 +2470,6 @@ def predict(
         # Data checks
         if Z.shape[0] != X.shape[0]:
             raise ValueError("X and Z must have the same number of rows")
-        if propensity is not None:
-            if propensity.shape[0] != X.shape[0]:
-                raise ValueError("X and propensity must have the same number of rows")
-        else:
-            if self.propensity_covariate != "none":
-                if not self.internal_propensity_model:
-                    raise ValueError(
-                        "Propensity scores not provided, but no propensity model was trained during sampling"
-                    )
-                else:
-                    propensity = np.mean(
-                        self.bart_propensity_model.predict(X), axis=1, keepdims=True
-                    )
 
         # Covariate preprocessing
         if not self._covariate_preprocessor._check_is_fitted():
@@ -2489,6 +2491,22 @@ def predict(
                 covariates_processed = X
         else:
             covariates_processed = self._covariate_preprocessor.transform(X)
+        
+        # Handle propensities
+        if propensity is not None:
+            if propensity.shape[0] != X.shape[0]:
+                raise ValueError("X and propensity must have the same number of rows")
+        else:
+            if self.propensity_covariate != "none":
+                if not self.internal_propensity_model:
+                    raise ValueError(
+                        "Propensity scores not provided, but no propensity model was trained during sampling"
+                    )
+                else:
+                    internal_propensity_preds = self.bart_propensity_model.predict(covariates_processed)
+                    propensity = np.mean(
+                        internal_propensity_preds['y_hat'], axis=1, keepdims=True
+                    )
 
         # Update covariates to include propensities if requested
         if self.propensity_covariate == "none":
@@ -2543,13 +2561,13 @@ def predict(
 
         # Return result matrices as a tuple
         if self.has_rfx and self.include_variance_forest:
-            return (tau_x, mu_x, rfx_preds, yhat_x, sigma2_x)
+            return {"mu_hat": mu_x, "tau_hat": tau_x, "y_hat": yhat_x, "rfx_predictions": rfx_preds, "variance_forest_predictions": sigma2_x}
         elif not self.has_rfx and self.include_variance_forest:
-            return (tau_x, mu_x, yhat_x, sigma2_x)
+            return {"mu_hat": mu_x, "tau_hat": tau_x, "y_hat": yhat_x, "rfx_predictions": None, "variance_forest_predictions": sigma2_x}
         elif self.has_rfx and not self.include_variance_forest:
-            return (tau_x, mu_x, rfx_preds, yhat_x)
+            return {"mu_hat": mu_x, "tau_hat": tau_x, "y_hat": yhat_x, "rfx_predictions": rfx_preds, "variance_forest_predictions": None}
         else:
-            return (tau_x, mu_x, yhat_x)
+            return {"mu_hat": mu_x, "tau_hat": tau_x, "y_hat": yhat_x, "rfx_predictions": None, "variance_forest_predictions": None}
 
     def to_json(self) -> str:
         """
diff --git a/test/python/test_bart.py b/test/python/test_bart.py
diff --git a/test/python/test_bcf.py b/test/python/test_bcf.py
diff --git a/test/python/test_json.py b/test/python/test_json.py
diff --git a/tools/perf/bart_profiling_script.py b/tools/perf/bart_profiling_script.py

Original file line number	Diff line number	Diff line change
`@@ -1276,15 +1276,23 @@ predict.bartmodel <- function(object, X, leaf_basis = NULL, rfx_group_ids = NULL`
`1276`	`1276`	`result <- list()`
`1277`	`1277`	`if ((object$model_params$has_rfx) \|\| (object$model_params$include_mean_forest)) {`
`1278`	`1278`	`result[["y_hat"]] = y_hat`
	`1279`	`+ } else {`
	`1280`	`+ result[["y_hat"]] <- NULL`
`1279`	`1281`	`}`
`1280`	`1282`	`if (object$model_params$include_mean_forest) {`
`1281`	`1283`	`result[["mean_forest_predictions"]] = mean_forest_predictions`
	`1284`	`+ } else {`
	`1285`	`+ result[["mean_forest_predictions"]] <- NULL`
`1282`	`1286`	`}`
`1283`	`1287`	`if (object$model_params$has_rfx) {`
`1284`	`1288`	`result[["rfx_predictions"]] = rfx_predictions`
	`1289`	`+ } else {`
	`1290`	`+ result[["rfx_predictions"]] <- NULL`
`1285`	`1291`	`}`
`1286`	`1292`	`if (object$model_params$include_variance_forest) {`
`1287`	`1293`	`result[["variance_forest_predictions"]] = variance_forest_predictions`
	`1294`	`+ } else {`
	`1295`	`+ result[["variance_forest_predictions"]] <- NULL`
`1288`	`1296`	`}`
`1289`	`1297`	`return(result)`
`1290`	`1298`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1800,10 +1800,14 @@ predict.bcfmodel <- function(object, X, Z, propensity = NULL, rfx_group_ids = NU`
`1800`	`1800`	`"y_hat" = y_hat`
`1801`	`1801`	`)`
`1802`	`1802`	`if (object$model_params$has_rfx) {`
`1803`		`- result[["rfx_predictions"]] = rfx_predictions`
	`1803`	`+ result[["rfx_predictions"]] <- rfx_predictions`
	`1804`	`+ } else {`
	`1805`	`+ result[["rfx_predictions"]] <- NULL`
`1804`	`1806`	`}`
`1805`	`1807`	`if (object$model_params$include_variance_forest) {`
`1806`		`- result[["variance_forest_predictions"]] = variance_forest_predictions`
	`1808`	`+ result[["variance_forest_predictions"]] <- variance_forest_predictions`
	`1809`	`+ } else {`
	`1810`	`+ result[["variance_forest_predictions"]] <- NULL`
`1807`	`1811`	`}`
`1808`	`1812`	`return(result)`
`1809`	`1813`	`}`
Original file line number	Diff line number	Diff line change
`@@ -145,7 +145,8 @@ def outcome_mean(X, W):`
`145`	`145`	`)`
`146`	`146`
`147`	`147`	`# Inspect the model outputs`
`148`		`- y_hat_mcmc = combined_bart.predict(X_test, basis_test)`
	`148`	`+ bart_preds = combined_bart.predict(X_test, basis_test)`
	`149`	`+ y_hat_mcmc = bart_preds['y_hat']`
`149`	`150`	`y_avg_mcmc = np.squeeze(y_hat_mcmc).mean(axis=1, keepdims=True)`
`150`	`151`	`y_df = pd.DataFrame(`
`151`	`152`	`np.concatenate((y_avg_mcmc, np.expand_dims(y_test, axis=1)), axis=1),`