Updated compute_contrast in Python and corrected other R package issues

Author: andrewherren

R/posterior_transformation.R

@@ -751,7 +751,7 @@ sample_bart_posterior_predictive <- function(
   # Compute posterior samples
   bart_preds <- predict(
     model_object,
-    covariates = covariates,
+    X = covariates,
     leaf_basis = basis,
     rfx_group_ids = rfx_group_ids,
     rfx_basis = rfx_basis,
@@ -1188,7 +1188,7 @@ compute_bart_posterior_interval <- function(
   # Compute posterior matrices for the requested model terms
   predictions <- predict(
     model_object,
-    covariates = covariates,
+    X = covariates,
     leaf_basis = basis,
     rfx_group_ids = rfx_group_ids,
     rfx_basis = rfx_basis,

R/utils.R

@@ -1092,33 +1092,3 @@ expand_dims_2d_diag <- function(input, output_size) {
   }
   return(output)
 }
-
-
-gfr_tie_checks <- function(covariates) {
-  num_vars <- ncol(covariates)
-  for (j in 1:num_vars) {
-    x_j <- covariates[, j]
-    if (has_few_unique_values(x_j)) {
-      warning_message <- paste0(
-        "Covariate column ",
-        j,
-        " has relatively few unique values. ",
-        "This may lead to tied values when sampling split points in BART/BCF, ",
-        "which can cause errors during model fitting. ",
-        "Consider adding small amounts of noise to this variable to break ties."
-      )
-      warning(warning_message)
-    }
-  }
-}
-
-
-has_few_unique_values <- function(
-  x,
-  count_threshold = 15
-) {
-  x_unique <- unique(x)
-  num_unique_values <- length(unique_values)
-  unique_to_total_count_ratio <- num_unique_values / length(x)
-  return(num_unique_values <= threshold)
-}

demo/debug/bart_contrast_debug.py

@@ -55,10 +55,10 @@
 
 # Compute contrast posterior
 contrast_posterior_test = bart_model.compute_contrast(
-    covariates_0=X_test,
-    covariates_1=X_test,
-    basis_0=np.zeros((n_test, 1)),
-    basis_1=np.ones((n_test, 1)),
+    X_0=X_test,
+    X_1=X_test,
+    leaf_basis_0=np.zeros((n_test, 1)),
+    leaf_basis_1=np.ones((n_test, 1)),
     type="posterior",
     scale="linear",
 )
@@ -143,10 +143,10 @@
 
 # Compute contrast posterior
 contrast_posterior_test = bart_model.compute_contrast(
-    covariates_0=X_test,
-    covariates_1=X_test,
-    basis_0=np.zeros((n_test, 1)),
-    basis_1=np.ones((n_test, 1)),
+    X_0=X_test,
+    X_1=X_test,
+    leaf_basis_0=np.zeros((n_test, 1)),
+    leaf_basis_1=np.ones((n_test, 1)),
     rfx_group_ids_0=group_ids_test,
     rfx_group_ids_1=group_ids_test,
     rfx_basis_0=rfx_basis_test,

demo/debug/probit_bart_rfx_debug.py

@@ -72,10 +72,10 @@
 
 # Compute contrast posterior
 contrast_posterior_test = bart_model.compute_contrast(
-    covariates_0=X_test,
-    covariates_1=X_test,
-    basis_0=np.zeros((n_test, 1)),
-    basis_1=np.ones((n_test, 1)),
+    X_0=X_test,
+    X_1=X_test,
+    leaf_basis_0=np.zeros((n_test, 1)),
+    leaf_basis_1=np.ones((n_test, 1)),
     rfx_group_ids_0=group_ids_test,
     rfx_group_ids_1=group_ids_test,
     rfx_basis_0=rfx_basis_test,

stochtree/bart.py

@@ -2046,10 +2046,10 @@ def predict(
 
     def compute_contrast(
         self,
-        covariates_0: Union[np.array, pd.DataFrame],
-        covariates_1: Union[np.array, pd.DataFrame],
-        basis_0: np.array = None,
-        basis_1: np.array = None,
+        X_0: Union[np.array, pd.DataFrame],
+        X_1: Union[np.array, pd.DataFrame],
+        leaf_basis_0: np.array = None,
+        leaf_basis_1: np.array = None,
         rfx_group_ids_0: np.array = None,
         rfx_group_ids_1: np.array = None,
         rfx_basis_0: np.array = None,
@@ -2068,13 +2068,13 @@ def compute_contrast(
 
         Parameters
         ----------
-        covariates_0 : np.array or pd.DataFrame
+        X_0 : np.array or pd.DataFrame
             Covariates used for prediction in the "control" case. Must be a numpy array or dataframe.
-        covariates_1 : np.array or pd.DataFrame
+        X_1 : np.array or pd.DataFrame
             Covariates used for prediction in the "treatment" case. Must be a numpy array or dataframe.
-        basis_0 : np.array, optional
+        leaf_basis_0 : np.array, optional
             Bases used for prediction in the "control" case (by e.g. dot product with leaf values).
-        basis_1 : np.array, optional
+        leaf_basis_1 : np.array, optional
             Bases used for prediction in the "treatment" case (by e.g. dot product with leaf values).
         rfx_group_ids_0 : np.array, optional
             Test set group labels used for prediction from an additive random effects model in the "control" case.
@@ -2135,33 +2135,33 @@ def compute_contrast(
             raise NotSampledError(msg)
 
         # Data checks
-        if not isinstance(covariates_0, pd.DataFrame) and not isinstance(
-            covariates_0, np.ndarray
+        if not isinstance(X_0, pd.DataFrame) and not isinstance(
+            X_0, np.ndarray
         ):
-            raise ValueError("covariates_0 must be a pandas dataframe or numpy array")
-        if not isinstance(covariates_1, pd.DataFrame) and not isinstance(
-            covariates_1, np.ndarray
+            raise ValueError("X_0 must be a pandas dataframe or numpy array")
+        if not isinstance(X_1, pd.DataFrame) and not isinstance(
+            X_1, np.ndarray
         ):
-            raise ValueError("covariates_1 must be a pandas dataframe or numpy array")
-        if basis_0 is not None:
-            if not isinstance(basis_0, np.ndarray):
-                raise ValueError("basis_0 must be a numpy array")
-            if basis_0.shape[0] != covariates_0.shape[0]:
+            raise ValueError("X_1 must be a pandas dataframe or numpy array")
+        if leaf_basis_0 is not None:
+            if not isinstance(leaf_basis_0, np.ndarray):
+                raise ValueError("leaf_basis_0 must be a numpy array")
+            if leaf_basis_0.shape[0] != X_0.shape[0]:
                 raise ValueError(
-                    "covariates_0 and basis_0 must have the same number of rows"
+                    "X_0 and leaf_basis_0 must have the same number of rows"
                 )
-        if basis_1 is not None:
-            if not isinstance(basis_1, np.ndarray):
-                raise ValueError("basis_1 must be a numpy array")
-            if basis_1.shape[0] != covariates_1.shape[0]:
+        if leaf_basis_1 is not None:
+            if not isinstance(leaf_basis_1, np.ndarray):
+                raise ValueError("leaf_basis_1 must be a numpy array")
+            if leaf_basis_1.shape[0] != X_1.shape[0]:
                 raise ValueError(
-                    "covariates_1 and basis_1 must have the same number of rows"
+                    "X_1 and leaf_basis_1 must have the same number of rows"
                 )
 
         # Predict for the control arm
         control_preds = self.predict(
-            covariates=covariates_0,
-            basis=basis_0,
+            X=X_0,
+            leaf_basis=leaf_basis_0,
             rfx_group_ids=rfx_group_ids_0,
             rfx_basis=rfx_basis_0,
             type="posterior",
@@ -2171,8 +2171,8 @@ def compute_contrast(
 
         # Predict for the treatment arm
         treatment_preds = self.predict(
-            covariates=covariates_1,
-            basis=basis_1,
+            X=X_1,
+            leaf_basis=leaf_basis_1,
             rfx_group_ids=rfx_group_ids_1,
             rfx_basis=rfx_basis_1,
             type="posterior",