Updated posterior predictive method / function in R and Python

Author: andrewherren

R/posterior_transformation.R

@@ -659,8 +659,8 @@ sample_bcf_posterior_predictive <- function(
 #' Sample from the posterior predictive distribution for outcomes modeled by BART
 #'
 #' @param model_object A fitted BART model object of class `bartmodel`.
-#' @param covariates A matrix or data frame of covariates. Required if the BART model depends on covariates (e.g., contains a mean or variance forest).
-#' @param basis A matrix of bases for mean forest models with regression defined in the leaves. Required for "leaf regression" models.
+#' @param X A matrix or data frame of covariates. Required if the BART model depends on covariates (e.g., contains a mean or variance forest).
+#' @param leaf_basis A matrix of bases for mean forest models with regression defined in the leaves. Required for "leaf regression" models.
 #' @param rfx_group_ids A vector of group IDs for random effects model. Required if the BART model includes random effects.
 #' @param rfx_basis A matrix of bases for random effects model. Required if the BART model includes random effects.
 #' @param num_draws_per_sample The number of posterior predictive samples to draw for each posterior sample. Defaults to a heuristic based on the number of samples in a BART model (i.e. if the BART model has >1000 draws, we use 1 draw from the likelihood per sample, otherwise we upsample to ensure intervals are based on at least 1000 posterior predictive draws).
@@ -675,12 +675,12 @@ sample_bcf_posterior_predictive <- function(
 #' y <- 2 * X[,1] + rnorm(n)
 #' bart_model <- bart(y_train = y, X_train = X)
 #' ppd_samples <- sample_bart_posterior_predictive(
-#'   model_object = bart_model, covariates = X
+#'   model_object = bart_model, X = X
 #' )
 sample_bart_posterior_predictive <- function(
   model_object,
-  covariates = NULL,
-  basis = NULL,
+  X = NULL,
+  leaf_basis = NULL,
   rfx_group_ids = NULL,
   rfx_basis = NULL,
   num_draws_per_sample = NULL
@@ -694,32 +694,32 @@ sample_bart_posterior_predictive <- function(
   # Check that all the necessary inputs were provided for interval computation
   needs_covariates <- model_object$model_params$include_mean_forest
   if (needs_covariates) {
-    if (is.null(covariates)) {
+    if (is.null(X)) {
       stop(
-        "'covariates' must be provided in order to compute the requested intervals"
+        "'X' must be provided in order to compute the requested intervals"
       )
     }
-    if (!is.matrix(covariates) && !is.data.frame(covariates)) {
-      stop("'covariates' must be a matrix or data frame")
+    if (!is.matrix(X) && !is.data.frame(X)) {
+      stop("'X' must be a matrix or data frame")
     }
   }
   needs_basis <- needs_covariates && model_object$model_params$has_basis
   if (needs_basis) {
-    if (is.null(basis)) {
+    if (is.null(leaf_basis)) {
       stop(
-        "'basis' must be provided in order to compute the requested intervals"
+        "'leaf_basis' must be provided in order to compute the requested intervals"
       )
     }
-    if (!is.matrix(basis)) {
-      stop("'basis' must be a matrix")
+    if (!is.matrix(leaf_basis)) {
+      stop("'leaf_basis' must be a matrix")
     }
-    if (is.matrix(basis)) {
-      if (nrow(basis) != nrow(covariates)) {
-        stop("'basis' must have the same number of rows as 'covariates'")
+    if (is.matrix(leaf_basis)) {
+      if (nrow(leaf_basis) != nrow(X)) {
+        stop("'leaf_basis' must have the same number of rows as 'X'")
       }
     } else {
-      if (length(basis) != nrow(covariates)) {
-        stop("'basis' must have the same number of elements as 'covariates'")
+      if (length(leaf_basis) != nrow(X)) {
+        stop("'leaf_basis' must have the same number of elements as 'X'")
       }
     }
   }
@@ -730,9 +730,9 @@ sample_bart_posterior_predictive <- function(
         "'rfx_group_ids' must be provided in order to compute the requested intervals"
       )
     }
-    if (length(rfx_group_ids) != nrow(covariates)) {
+    if (length(rfx_group_ids) != nrow(X)) {
       stop(
-        "'rfx_group_ids' must have the same length as the number of rows in 'covariates'"
+        "'rfx_group_ids' must have the same length as the number of rows in 'X'"
       )
     }
     if (is.null(rfx_basis)) {
@@ -743,16 +743,16 @@ sample_bart_posterior_predictive <- function(
     if (!is.matrix(rfx_basis)) {
       stop("'rfx_basis' must be a matrix")
     }
-    if (nrow(rfx_basis) != nrow(covariates)) {
-      stop("'rfx_basis' must have the same number of rows as 'covariates'")
+    if (nrow(rfx_basis) != nrow(X)) {
+      stop("'rfx_basis' must have the same number of rows as 'X'")
     }
   }
 
   # Compute posterior samples
   bart_preds <- predict(
     model_object,
-    X = covariates,
-    leaf_basis = basis,
+    X = X,
+    leaf_basis = leaf_basis,
     rfx_group_ids = rfx_group_ids,
     rfx_basis = rfx_basis,
     type = "posterior",
@@ -766,7 +766,7 @@ sample_bart_posterior_predictive <- function(
   has_variance_forest <- model_object$model_params$include_variance_forest
   samples_global_variance <- model_object$model_params$sample_sigma2_global
   num_posterior_draws <- model_object$model_params$num_samples
-  num_observations <- nrow(covariates)
+  num_observations <- nrow(X)
   if (has_mean_term) {
     ppd_mean <- bart_preds$y_hat
   } else {

man/sample_bart_posterior_predictive.Rd

@@ -2338,8 +2338,8 @@ def compute_posterior_interval(
 
     def sample_posterior_predictive(
         self,
-        covariates: np.array = None,
-        basis: np.array = None,
+        X: np.array = None,
+        leaf_basis: np.array = None,
         rfx_group_ids: np.array = None,
         rfx_basis: np.array = None,
         num_draws_per_sample: int = None,
@@ -2349,9 +2349,9 @@ def sample_posterior_predictive(
 
         Parameters
         ----------
-        covariates : np.array, optional
+        X : np.array, optional
             An array or data frame of covariates at which to compute the intervals. Required if the BART model depends on covariates (e.g., contains a mean or variance forest).
-        basis : np.array, optional
+        leaf_basis : np.array, optional
             An array of basis function evaluations for mean forest models with regression defined in the leaves. Required for "leaf regression" models.
         rfx_group_ids : np.array, optional
             An array of group IDs for random effects. Required if the BART model includes random effects.
@@ -2375,25 +2375,25 @@ def sample_posterior_predictive(
         # Check that all the necessary inputs were provided for interval computation
         needs_covariates = self.include_mean_forest
         if needs_covariates:
-            if covariates is None:
+            if X is None:
                 raise ValueError(
-                    "'covariates' must be provided in order to compute the requested intervals"
+                    "'X' must be provided in order to compute the requested intervals"
                 )
-            if not isinstance(covariates, np.ndarray) and not isinstance(
-                covariates, pd.DataFrame
+            if not isinstance(X, np.ndarray) and not isinstance(
+                X, pd.DataFrame
             ):
-                raise ValueError("'covariates' must be a matrix or data frame")
+                raise ValueError("'X' must be a matrix or data frame")
         needs_basis = needs_covariates and self.has_basis
         if needs_basis:
-            if basis is None:
+            if leaf_basis is None:
                 raise ValueError(
-                    "'basis' must be provided in order to compute the requested intervals"
+                    "'leaf_basis' must be provided in order to compute the requested intervals"
                 )
-            if not isinstance(basis, np.ndarray):
-                raise ValueError("'basis' must be a numpy array")
-            if basis.shape[0] != covariates.shape[0]:
+            if not isinstance(leaf_basis, np.ndarray):
+                raise ValueError("'leaf_basis' must be a numpy array")
+            if leaf_basis.shape[0] != X.shape[0]:
                 raise ValueError(
-                    "'basis' must have the same number of rows as 'covariates'"
+                    "'leaf_basis' must have the same number of rows as 'X'"
                 )
         needs_rfx_data = self.has_rfx
         if needs_rfx_data:
@@ -2403,25 +2403,25 @@ def sample_posterior_predictive(
                 )
             if not isinstance(rfx_group_ids, np.ndarray):
                 raise ValueError("'rfx_group_ids' must be a numpy array")
-            if rfx_group_ids.shape[0] != covariates.shape[0]:
+            if rfx_group_ids.shape[0] != X.shape[0]:
                 raise ValueError(
-                    "'rfx_group_ids' must have the same length as the number of rows in 'covariates'"
+                    "'rfx_group_ids' must have the same length as the number of rows in 'X'"
                 )
             if rfx_basis is None:
                 raise ValueError(
                     "'rfx_basis' must be provided in order to compute the requested intervals"
                 )
             if not isinstance(rfx_basis, np.ndarray):
                 raise ValueError("'rfx_basis' must be a numpy array")
-            if rfx_basis.shape[0] != covariates.shape[0]:
+            if rfx_basis.shape[0] != X.shape[0]:
                 raise ValueError(
-                    "'rfx_basis' must have the same number of rows as 'covariates'"
+                    "'rfx_basis' must have the same number of rows as 'X'"
                 )
 
         # Compute posterior predictive samples
         bart_preds = self.predict(
-            covariates=covariates,
-            basis=basis,
+            X=X,
+            leaf_basis=leaf_basis,
             rfx_group_ids=rfx_group_ids,
             rfx_basis=rfx_basis,
             type="posterior",
@@ -2433,7 +2433,7 @@ def sample_posterior_predictive(
         has_variance_forest = self.include_variance_forest
         samples_global_variance = self.sample_sigma2_global
         num_posterior_draws = self.num_samples
-        num_observations = covariates.shape[0]
+        num_observations = X.shape[0]
         if has_mean_term:
             ppd_mean = bart_preds["y_hat"]
         else:

stochtree/bart.py

@@ -67,7 +67,7 @@ y_hat_intervals <- compute_bart_posterior_interval(
 
 pred_intervals <- sample_bart_posterior_predictive(
   model_object = bart_model,
-  covariates = X_test,
+  X = X_test,
   level = 0.95
 )
 
@@ -169,7 +169,7 @@ lines(y_hat_prob_intervals$upper[sort_inds])
 # Draw from posterior predictive for covariates in the test set
 ppd_samples <- sample_bart_posterior_predictive(
   model_object = bart_model,
-  covariates = X_test,
+  X = X_test,
   num_draws = 10
 )