Updated tests and BART and BCF R functions

andrewherren · andrewherren · commit 9679068fd1fe · 2025-11-20T00:50:36.000-06:00
diff --git a/R/bart.R b/R/bart.R
@@ -418,7 +418,11 @@ bart <- function(
   # Raise a warning if the data have ties and only GFR is being run
   if ((num_gfr > 0) && (num_mcmc == 0) && (num_burnin == 0)) {
     num_values <- nrow(X_train)
-    max_grid_size <- floor(num_values / cutpoint_grid_size)
+    max_grid_size <- ifelse(
+      num_values > cutpoint_grid_size,
+      floor(num_values / cutpoint_grid_size),
+      1
+    )
     covs_warning_1 <- NULL
     covs_warning_2 <- NULL
     covs_warning_3 <- NULL
diff --git a/R/bcf.R b/R/bcf.R
@@ -522,7 +522,11 @@ bcf <- function(
   # Raise a warning if the data have ties and only GFR is being run
   if ((num_gfr > 0) && (num_mcmc == 0) && (num_burnin == 0)) {
     num_values <- nrow(X_train)
-    max_grid_size <- floor(num_values / cutpoint_grid_size)
+    max_grid_size <- ifelse(
+      num_values > cutpoint_grid_size,
+      floor(num_values / cutpoint_grid_size),
+      1
+    )
     covs_warning_1 <- NULL
     covs_warning_2 <- NULL
     covs_warning_3 <- NULL
diff --git a/test/R/testthat/test-bart.R b/test/R/testthat/test-bart.R
@@ -312,6 +312,19 @@ test_that("Warmstart BART", {
   # Run a new BART chain from the existing (X)BART model
   general_param_list <- list(num_chains = 3, keep_every = 5)
   expect_no_error(
+    bart_model <- bart(
+      X_train = X_train,
+      y_train = y_train,
+      X_test = X_test,
+      num_gfr = 0,
+      num_burnin = 10,
+      num_mcmc = 10,
+      previous_model_json = bart_model_json_string,
+      previous_model_warmstart_sample_num = 10,
+      general_params = general_param_list
+    )
+  )
+  expect_warning(
     bart_model <- bart(
       X_train = X_train,
       y_train = y_train,
@@ -376,6 +389,23 @@ test_that("Warmstart BART", {
   # Run a new BART chain from the existing (X)BART model
   general_param_list <- list(num_chains = 4, keep_every = 5)
   expect_no_error(
+    bart_model <- bart(
+      X_train = X_train,
+      y_train = y_train,
+      X_test = X_test,
+      rfx_group_ids_train = rfx_group_ids_train,
+      rfx_group_ids_test = rfx_group_ids_test,
+      rfx_basis_train = rfx_basis_train,
+      rfx_basis_test = rfx_basis_test,
+      num_gfr = 0,
+      num_burnin = 10,
+      num_mcmc = 10,
+      previous_model_json = bart_model_json_string,
+      previous_model_warmstart_sample_num = 10,
+      general_params = general_param_list
+    )
+  )
+  expect_warning(
     bart_model <- bart(
       X_train = X_train,
       y_train = y_train,
diff --git a/test/R/testthat/test-bcf.R b/test/R/testthat/test-bcf.R
@@ -375,6 +375,23 @@ test_that("Warmstart BCF", {
   # Run a new BCF chain from the existing (X)BCF model
   general_param_list <- list(num_chains = 3, keep_every = 5)
   expect_no_error(
+    bcf_model <- bcf(
+      X_train = X_train,
+      y_train = y_train,
+      Z_train = Z_train,
+      propensity_train = pi_train,
+      X_test = X_test,
+      Z_test = Z_test,
+      propensity_test = pi_test,
+      num_gfr = 0,
+      num_burnin = 10,
+      num_mcmc = 10,
+      previous_model_json = bcf_model_json_string,
+      previous_model_warmstart_sample_num = 10,
+      general_params = general_param_list
+    )
+  )
+  expect_warning(
     bcf_model <- bcf(
       X_train = X_train,
       y_train = y_train,
@@ -482,7 +499,7 @@ test_that("Warmstart BCF", {
       num_burnin = 10,
       num_mcmc = 10,
       previous_model_json = bcf_model_json_string,
-      previous_model_warmstart_sample_num = 1,
+      previous_model_warmstart_sample_num = 10,
       general_params = general_param_list
     )
   )
diff --git a/test/python/test_bart.py b/test/python/test_bart.py
@@ -83,7 +83,7 @@ def outcome_mean(X):
         bart_model_3.from_json_string_list(bart_models_json)
 
         # Assertions
-        bart_preds_combined = bart_model_3.predict(covariates=X_train)
+        bart_preds_combined = bart_model_3.predict(X=X_train)
         y_hat_train_combined = bart_preds_combined["y_hat"]
         assert y_hat_train_combined.shape == (n_train, num_mcmc * 2)
         np.testing.assert_allclose(
@@ -190,7 +190,7 @@ def outcome_mean(X, W):
 
         # Assertions
         bart_preds_combined = bart_model_3.predict(
-            covariates=X_train, basis=basis_train
+            X=X_train, leaf_basis=basis_train
         )
         y_hat_train_combined = bart_preds_combined["y_hat"]
         assert y_hat_train_combined.shape == (n_train, num_mcmc * 2)
@@ -298,7 +298,7 @@ def outcome_mean(X, W):
 
         # Assertions
         bart_preds_combined = bart_model_3.predict(
-            covariates=X_train, basis=basis_train
+            X=X_train, leaf_basis=basis_train
         )
         y_hat_train_combined = bart_preds_combined["y_hat"]
         assert y_hat_train_combined.shape == (n_train, num_mcmc * 2)
@@ -410,7 +410,7 @@ def conditional_stddev(X):
         bart_model_3.from_json_string_list(bart_models_json)
 
         # Assertions
-        bart_preds_combined = bart_model_3.predict(covariates=X_train)
+        bart_preds_combined = bart_model_3.predict(X=X_train)
         y_hat_train_combined, sigma2_x_train_combined = (
             bart_preds_combined["y_hat"],
             bart_preds_combined["variance_forest_predictions"],
@@ -545,7 +545,7 @@ def conditional_stddev(X):
 
         # Assertions
         bart_preds_combined = bart_model_3.predict(
-            covariates=X_train, basis=basis_train
+            X=X_train, leaf_basis=basis_train
         )
         y_hat_train_combined = bart_preds_combined["y_hat"]
         assert y_hat_train_combined.shape == (n_train, num_mcmc * 2)
@@ -670,7 +670,7 @@ def conditional_stddev(X):
 
         # Assertions
         bart_preds_combined = bart_model_3.predict(
-            covariates=X_train, basis=basis_train
+            X=X_train, leaf_basis=basis_train
         )
         y_hat_train_combined = bart_preds_combined["y_hat"]
         assert y_hat_train_combined.shape == (n_train, num_mcmc * 2)
@@ -825,7 +825,7 @@ def rfx_term(group_labels, basis):
 
         # Assertions
         bart_preds_combined = bart_model_3.predict(
-            covariates=X_train,
+            X=X_train,
             rfx_group_ids=group_labels_train,
             rfx_basis=rfx_basis_train,
         )
@@ -998,8 +998,8 @@ def conditional_stddev(X):
 
         # Assertions
         bart_preds_combined = bart_model_3.predict(
-            covariates=X_train,
-            basis=basis_train,
+            X=X_train,
+            leaf_basis=basis_train,
             rfx_group_ids=group_labels_train,
             rfx_basis=rfx_basis_train,
         )
@@ -1196,8 +1196,8 @@ def conditional_stddev(X):
             random_effects_params=rfx_params,
         )
         preds = bart_model_4.predict(
-            covariates=X_test,
-            basis=basis_test,
+            X=X_test,
+            leaf_basis=basis_test,
             rfx_group_ids=group_labels_test,
             type="posterior",
             terms="rfx",
diff --git a/test/python/test_predict.py b/test/python/test_predict.py
@@ -221,12 +221,12 @@ def test_bart_prediction(self):
         )
 
         # Check that the default predict method returns a dictionary
-        pred = bart_model.predict(covariates=X_test)
+        pred = bart_model.predict(X=X_test)
         y_hat_posterior_test = pred["y_hat"]
         assert y_hat_posterior_test.shape == (20, 10)
 
         # Check that the pre-aggregated predictions match with those computed by np.mean
-        pred_mean = bart_model.predict(covariates=X_test, type="mean")
+        pred_mean = bart_model.predict(X=X_test, type="mean")
         y_hat_mean_test = pred_mean["y_hat"]
         np.testing.assert_almost_equal(
             y_hat_mean_test, np.mean(y_hat_posterior_test, axis=1)
@@ -245,14 +245,14 @@ def test_bart_prediction(self):
         )
 
         # Check that the default predict method returns a dictionary
-        pred = het_bart_model.predict(covariates=X_test)
+        pred = het_bart_model.predict(X=X_test)
         y_hat_posterior_test = pred["y_hat"]
         sigma2_hat_posterior_test = pred["variance_forest_predictions"]
         assert y_hat_posterior_test.shape == (20, 10)
         assert sigma2_hat_posterior_test.shape == (20, 10)
 
         # Check that the pre-aggregated predictions match with those computed by np.mean
-        pred_mean = het_bart_model.predict(covariates=X_test, type="mean")
+        pred_mean = het_bart_model.predict(X=X_test, type="mean")
         y_hat_mean_test = pred_mean["y_hat"]
         sigma2_hat_mean_test = pred_mean["variance_forest_predictions"]
         np.testing.assert_almost_equal(
@@ -265,10 +265,10 @@ def test_bart_prediction(self):
         # Check that the "single-term" pre-aggregated predictions
         # match those computed by pre-aggregated predictions returned in a dictionary
         y_hat_mean_test_single_term = het_bart_model.predict(
-            covariates=X_test, type="mean", terms="y_hat"
+            X=X_test, type="mean", terms="y_hat"
         )
         sigma2_hat_mean_test_single_term = het_bart_model.predict(
-            covariates=X_test, type="mean", terms="variance_forest"
+            X=X_test, type="mean", terms="variance_forest"
         )
         np.testing.assert_almost_equal(y_hat_mean_test, y_hat_mean_test_single_term)
         np.testing.assert_almost_equal(
