Harmonized BART sampler arguments between R and Python (and wherever they are called)

Author: andrewherren

R/bart.R

@@ -1924,7 +1924,7 @@ bart <- function(
 #' Predict from a sampled BART model on new data
 #'
 #' @param object Object of type `bart` containing draws of a regression forest and associated sampling outputs.
-#' @param covariates Covariates used to determine tree leaf predictions for each observation. Must be passed as a matrix or dataframe.
+#' @param X Covariates used to determine tree leaf predictions for each observation. Must be passed as a matrix or dataframe.
 #' @param leaf_basis (Optional) Bases used for prediction (by e.g. dot product with leaf values). Default: `NULL`.
 #' @param rfx_group_ids (Optional) Test set group labels used for an additive random effects model.
 #' We do not currently support (but plan to in the near future), test set evaluation for group labels
@@ -1964,7 +1964,7 @@ bart <- function(
 #' y_hat_test <- predict(bart_model, X_test)$y_hat
 predict.bartmodel <- function(
   object,
-  covariates,
+  X,
   leaf_basis = NULL,
   rfx_group_ids = NULL,
   rfx_basis = NULL,
@@ -2047,8 +2047,8 @@ predict.bartmodel <- function(
   }
 
   # Check that covariates are matrix or data frame
-  if ((!is.data.frame(covariates)) && (!is.matrix(covariates))) {
-    stop("covariates must be a matrix or dataframe")
+  if ((!is.data.frame(X)) && (!is.matrix(X))) {
+    stop("X must be a matrix or dataframe")
   }
 
   # Convert all input data to matrices if not already converted
@@ -2063,12 +2063,12 @@ predict.bartmodel <- function(
   if ((object$model_params$requires_basis) && (is.null(leaf_basis))) {
     stop("Basis (leaf_basis) must be provided for this model")
   }
-  if ((!is.null(leaf_basis)) && (nrow(covariates) != nrow(leaf_basis))) {
-    stop("covariates and leaf_basis must have the same number of rows")
+  if ((!is.null(leaf_basis)) && (nrow(X) != nrow(leaf_basis))) {
+    stop("X and leaf_basis must have the same number of rows")
   }
-  if (object$model_params$num_covariates != ncol(covariates)) {
+  if (object$model_params$num_covariates != ncol(X)) {
     stop(
-      "covariates must contain the same number of columns as the BART model's training dataset"
+      "X must contain the same number of columns as the BART model's training dataset"
     )
   }
   if ((predict_rfx) && (is.null(rfx_group_ids))) {
@@ -2089,7 +2089,7 @@ predict.bartmodel <- function(
 
   # Preprocess covariates
   train_set_metadata <- object$train_set_metadata
-  covariates <- preprocessPredictionData(covariates, train_set_metadata)
+  X <- preprocessPredictionData(X, train_set_metadata)
 
   # Recode group IDs to integer vector (if passed as, for example, a vector of county names, etc...)
   has_rfx <- FALSE
@@ -2119,8 +2119,8 @@ predict.bartmodel <- function(
       # Only construct a basis if user-provided basis missing
       if (is.null(rfx_basis)) {
         rfx_basis <- matrix(
-          rep(1, nrow(covariates)),
-          nrow = nrow(covariates),
+          rep(1, nrow(X)),
+          nrow = nrow(X),
           ncol = 1
         )
       }
@@ -2129,9 +2129,9 @@ predict.bartmodel <- function(
 
   # Create prediction dataset
   if (!is.null(leaf_basis)) {
-    prediction_dataset <- createForestDataset(covariates, leaf_basis)
+    prediction_dataset <- createForestDataset(X, leaf_basis)
   } else {
-    prediction_dataset <- createForestDataset(covariates)
+    prediction_dataset <- createForestDataset(X)
   }
 
   # Compute variance forest predictions

demo/debug/bart_contrast_debug.py

@@ -65,15 +65,15 @@
 
 # Compute the same quantity via two predict calls
 y_hat_posterior_test_0 = bart_model.predict(
-    covariates=X_test,
-    basis=np.zeros((n_test, 1)),
+    X=X_test,
+    leaf_basis=np.zeros((n_test, 1)),
     type="posterior",
     terms="y_hat",
     scale="linear",
 )
 y_hat_posterior_test_1 = bart_model.predict(
-    covariates=X_test,
-    basis=np.ones((n_test, 1)),
+    X=X_test,
+    leaf_basis=np.ones((n_test, 1)),
     type="posterior",
     terms="y_hat",
     scale="linear",
@@ -157,17 +157,17 @@
 
 # Compute the same quantity via two predict calls
 y_hat_posterior_test_0 = bart_model.predict(
-    covariates=X_test,
-    basis=np.zeros((n_test, 1)),
+    X=X_test,
+    leaf_basis=np.zeros((n_test, 1)),
     rfx_group_ids=group_ids_test,
     rfx_basis=rfx_basis_test,
     type="posterior",
     terms="y_hat",
     scale="linear",
 )
 y_hat_posterior_test_1 = bart_model.predict(
-    covariates=X_test,
-    basis=np.ones((n_test, 1)),
+    X=X_test,
+    leaf_basis=np.ones((n_test, 1)),
     rfx_group_ids=group_ids_test,
     rfx_basis=rfx_basis_test,
     type="posterior",

demo/debug/bart_predict_debug.py

@@ -46,11 +46,11 @@
 )
 
 # # Check several predict approaches
-bart_preds = bart_model.predict(covariates=X_test)
-y_hat_posterior_test = bart_model.predict(covariates=X_test)["y_hat"]
-y_hat_mean_test = bart_model.predict(covariates=X_test, type="mean", terms=["y_hat"])
+bart_preds = bart_model.predict(X=X_test)
+y_hat_posterior_test = bart_model.predict(X=X_test)["y_hat"]
+y_hat_mean_test = bart_model.predict(X=X_test, type="mean", terms=["y_hat"])
 y_hat_test = bart_model.predict(
-    covariates=X_test, type="mean", terms=["rfx", "variance"]
+    X=X_test, type="mean", terms=["rfx", "variance"]
 )
 
 # Plot predicted versus actual

demo/debug/gfr_ties_debug.py

@@ -38,7 +38,7 @@
 )
 
 # Inspect the model fit
-y_hat_test = xbart_model.predict(X_test, type="mean", terms="y_hat")
+y_hat_test = xbart_model.predict(X=X_test, type="mean", terms="y_hat")
 plt.scatter(y_hat_test, y_test)
 plt.axline((0, 0), slope=1, color="red", linestyle="dashed", linewidth=1.5)
 plt.xlabel("Predicted Outcome Mean")
@@ -54,7 +54,7 @@
 )
 
 # Inspect the model fit
-y_hat_test = bart_model.predict(X_test, type="mean", terms="y_hat")
+y_hat_test = bart_model.predict(X=X_test, type="mean", terms="y_hat")
 plt.clf()
 plt.scatter(y_hat_test, y_test)
 plt.axline((0, 0), slope=1, color="red", linestyle="dashed", linewidth=1.5)
@@ -95,7 +95,7 @@
 )
 
 # Inspect the model fit
-y_hat_test = xbart_model.predict(X_test, type="mean", terms="y_hat")
+y_hat_test = xbart_model.predict(X=X_test, type="mean", terms="y_hat")
 plt.clf()
 plt.scatter(y_hat_test, y_test)
 plt.axline((0, 0), slope=1, color="red", linestyle="dashed", linewidth=1.5)
@@ -112,7 +112,7 @@
 )
 
 # Inspect the model fit
-y_hat_test = bart_model.predict(X_test, type="mean", terms="y_hat")
+y_hat_test = bart_model.predict(X=X_test, type="mean", terms="y_hat")
 plt.clf()
 plt.scatter(y_hat_test, y_test)
 plt.axline((0, 0), slope=1, color="red", linestyle="dashed", linewidth=1.5)

demo/debug/multi_chain.py

@@ -89,8 +89,8 @@ def outcome_mean(X, W):
 
 # Analyze model predictions collectively across all chains
 y_hat_test = bart_model.predict(
-  covariates = X_test,
-  basis = basis_test,
+  X = X_test,
+  leaf_basis = basis_test,
   type = "mean", 
   terms = "y_hat"
 )

demo/debug/multiple_initializations.py

@@ -118,14 +118,14 @@ def outcome_mean(X, W):
 )
 
 # Inspect the model outputs
-bart_preds_2 = bart_model_2.predict(X_test, basis_test)
+bart_preds_2 = bart_model_2.predict(X=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_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)
+bart_preds_3 = bart_model_3.predict(X=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)
-bart_preds_4 = bart_model_4.predict(X_test, basis_test)
+bart_preds_4 = bart_model_4.predict(X=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(

demo/debug/parallel_multi_chain.py

@@ -145,7 +145,7 @@ def outcome_mean(X, W):
     )
 
     # Inspect the model outputs
-    bart_preds = combined_bart.predict(X_test, basis_test)
+    bart_preds = combined_bart.predict(X=X_test, leaf_basis=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(

demo/debug/probit_bart_rfx_debug.py

@@ -86,17 +86,17 @@
 
 # Compute the same quantity via two predict calls
 y_hat_posterior_test_0 = bart_model.predict(
-    covariates=X_test,
-    basis=np.zeros((n_test, 1)),
+    X=X_test,
+    leaf_basis=np.zeros((n_test, 1)),
     rfx_group_ids=group_ids_test,
     rfx_basis=rfx_basis_test,
     type="posterior",
     terms="y_hat",
     scale="linear",
 )
 y_hat_posterior_test_1 = bart_model.predict(
-    covariates=X_test,
-    basis=np.ones((n_test, 1)),
+    X=X_test,
+    leaf_basis=np.ones((n_test, 1)),
     rfx_group_ids=group_ids_test,
     rfx_basis=rfx_basis_test,
     type="posterior",
@@ -111,8 +111,8 @@
 
 # Plot predicted versus actual outcome
 Z_hat_test = bart_model.predict(
-    covariates=X_test,
-    basis=W_test,
+    X=X_test,
+    leaf_basis=W_test,
     rfx_group_ids=group_ids_test,
     rfx_basis=rfx_basis_test,
     type="mean",

demo/debug/rfx_serialization.py

@@ -60,13 +60,13 @@ def rfx_mean(group_labels, basis):
                   rfx_basis_train=basis, num_gfr=10, num_mcmc=10)
 
 # Extract predictions from the sampler
-bart_preds_orig = bart_orig.predict(X, W, group_labels, basis)
+bart_preds_orig = bart_orig.predict(X=X, leaf_basis=W, rfx_group_ids=group_labels, rfx_basis=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)
-bart_preds_reloaded = bart_reloaded.predict(X, W, group_labels, basis)
+bart_preds_reloaded = bart_reloaded.predict(X=X, leaf_basis=W, rfx_group_ids=group_labels, rfx_basis=basis)
 y_hat_reloaded = bart_preds_reloaded['y_hat']
 np.testing.assert_almost_equal(y_hat_orig, y_hat_reloaded)

demo/notebooks/multi_chain.ipynb

@@ -161,8 +161,8 @@
    "outputs": [],
    "source": [
     "y_hat_test = bart_model.predict(\n",
-    "  covariates = X_test,\n",
-    "  basis = leaf_basis_test, \n",
+    "  X = X_test,\n",
+    "  leaf_basis = leaf_basis_test, \n",
     "  type = \"mean\", \n",
     "  terms = \"y_hat\"\n",
     ")\n",
@@ -321,8 +321,8 @@
    "outputs": [],
    "source": [
     "y_hat_test = bart_model.predict(\n",
-    "  covariates = X_test,\n",
-    "  basis = leaf_basis_test, \n",
+    "  X = X_test,\n",
+    "  leaf_basis = leaf_basis_test, \n",
     "  type = \"mean\", \n",
     "  terms = \"y_hat\"\n",
     ")\n",