Updated predict interface in R and Python and tests / demos that call it

Author: andrewherren

R/bart.R

@@ -1961,7 +1961,7 @@ bart <- function(
 #' y_train <- y[train_inds]
 #' bart_model <- bart(X_train = X_train, y_train = y_train,
 #'                    num_gfr = 10, num_burnin = 0, num_mcmc = 10)
-#' y_hat_test <- predict(bart_model, X_test)$y_hat
+#' y_hat_test <- predict(bart_model, X=X_test)$y_hat
 predict.bartmodel <- function(
   object,
   X,
@@ -2843,7 +2843,7 @@ createBARTModelFromJsonFile <- function(json_filename) {
 #'                    num_gfr = 10, num_burnin = 0, num_mcmc = 10)
 #' bart_json <- saveBARTModelToJsonString(bart_model)
 #' bart_model_roundtrip <- createBARTModelFromJsonString(bart_json)
-#' y_hat_mean_roundtrip <- rowMeans(predict(bart_model_roundtrip, X_train)$y_hat)
+#' y_hat_mean_roundtrip <- rowMeans(predict(bart_model_roundtrip, X=X_train)$y_hat)
 createBARTModelFromJsonString <- function(json_string) {
   # Load a `CppJson` object from string
   bart_json <- createCppJsonString(json_string)

R/kernel.R

@@ -129,7 +129,7 @@ computeForestLeafIndices <- function(
         propensity <- rowMeans(
           predict(
             model_object$bart_propensity_model,
-            covariates
+            X = covariates
           )$y_hat
         )
       }

man/createBARTModelFromJsonString.Rd

@@ -584,7 +584,7 @@ test_that("Random Effects BART", {
     )
     preds <- predict(
       bart_model,
-      covariates = X_test,
+      X = X_test,
       leaf_basis = W_test,
       rfx_group_ids = rfx_group_ids_test,
       type = "posterior",

man/predict.bartmodel.Rd

@@ -216,20 +216,20 @@ test_that("BART predictions with pre-summarization", {
   )
 
   # Check that the default predict method returns a list
-  pred <- predict(bart_model, X_test)
+  pred <- predict(bart_model, X = X_test)
   y_hat_posterior_test <- pred$y_hat
   expect_equal(dim(y_hat_posterior_test), c(20, 10))
 
   # Check that the pre-aggregated predictions match with those computed by rowMeans
-  pred_mean <- predict(bart_model, X_test, type = "mean")
+  pred_mean <- predict(bart_model, X = X_test, type = "mean")
   y_hat_mean_test <- pred_mean$y_hat
   expect_equal(y_hat_mean_test, rowMeans(y_hat_posterior_test))
 
   # Check that we warn and return a NULL when requesting terms that weren't fit
   expect_warning({
     pred_mean <- predict(
       bart_model,
-      X_test,
+      X = X_test,
       type = "mean",
       terms = c("rfx", "variance_forest")
     )
@@ -248,7 +248,7 @@ test_that("BART predictions with pre-summarization", {
   )
 
   # Check that the default predict method returns a list
-  pred <- predict(het_bart_model, X_test)
+  pred <- predict(het_bart_model, X = X_test)
   y_hat_posterior_test <- pred$y_hat
   sigma2_hat_posterior_test <- pred$variance_forest_predictions
 
@@ -257,7 +257,7 @@ test_that("BART predictions with pre-summarization", {
   expect_equal(dim(sigma2_hat_posterior_test), c(20, 10))
 
   # Check that the pre-aggregated predictions match with those computed by rowMeans
-  pred_mean <- predict(het_bart_model, X_test, type = "mean")
+  pred_mean <- predict(het_bart_model, X = X_test, type = "mean")
   y_hat_mean_test <- pred_mean$y_hat
   sigma2_hat_mean_test <- pred_mean$variance_forest_predictions
 
@@ -269,13 +269,13 @@ test_that("BART predictions with pre-summarization", {
   # match those computed by pre-aggregated predictions returned in a list
   y_hat_mean_test_single_term <- predict(
     het_bart_model,
-    X_test,
+    X = X_test,
     type = "mean",
     terms = "y_hat"
   )
   sigma2_hat_mean_test_single_term <- predict(
     het_bart_model,
-    X_test,
+    X = X_test,
     type = "mean",
     terms = "variance_forest"
   )

test/R/testthat/test-bart.R

@@ -34,7 +34,7 @@ test_that("BART Serialization", {
     num_mcmc = 10,
     general_params = general_param_list
   )
-  y_hat_orig <- rowMeans(predict(bart_model, X_test)$y_hat)
+  y_hat_orig <- rowMeans(predict(bart_model, X = X_test)$y_hat)
 
   # Save to JSON
   bart_json_string <- saveBARTModelToJsonString(bart_model)
@@ -43,7 +43,7 @@ test_that("BART Serialization", {
   bart_model_roundtrip <- createBARTModelFromJsonString(bart_json_string)
 
   # Predict from the roundtrip BART model
-  y_hat_reloaded <- rowMeans(predict(bart_model_roundtrip, X_test)$y_hat)
+  y_hat_reloaded <- rowMeans(predict(bart_model_roundtrip, X = X_test)$y_hat)
 
   # Assertion
   expect_equal(y_hat_orig, y_hat_reloaded)

test/R/testthat/test-predict.R

@@ -38,16 +38,16 @@ bart_model <- bart(
 )
 
 # Check several predict approaches
-y_hat_posterior_test <- predict(bart_model, X_test)$y_hat
+y_hat_posterior_test <- predict(bart_model, X = X_test)$y_hat
 y_hat_mean_test <- predict(
   bart_model,
-  X_test,
+  X = X_test,
   type = "mean",
   terms = c("y_hat")
 )
 y_hat_test <- predict(
   bart_model,
-  X_test,
+  X = X_test,
   type = "mean",
   terms = c("rfx", "variance")
 )
@@ -117,18 +117,18 @@ bart_model <- bart(
 # Predict on latent scale
 y_hat_post <- predict(
   object = bart_model,
+  X = X_test,
   type = "posterior",
   terms = c("y_hat"),
-  X = X_test,
   scale = "linear"
 )
 
 # Predict on probability scale
 y_hat_post_prob <- predict(
   object = bart_model,
+  X = X_test,
   type = "posterior",
   terms = c("y_hat"),
-  X = X_test,
   scale = "probability"
 )
 

test/R/testthat/test-serialization.R

@@ -14,35 +14,55 @@ num_trees <- 100
 n <- 500
 p_x <- 20
 snr <- 2
-X <- matrix(runif(n*p_x), ncol = p_x)
-f_XW <- sin(4*pi*X[,1]) + cos(4*pi*X[,2]) + sin(4*pi*X[,3]) +cos(4*pi*X[,4])
+X <- matrix(runif(n * p_x), ncol = p_x)
+f_XW <- sin(4 * pi * X[, 1]) +
+  cos(4 * pi * X[, 2]) +
+  sin(4 * pi * X[, 3]) +
+  cos(4 * pi * X[, 4])
 noise_sd <- sd(f_XW) / snr
-y <- f_XW + rnorm(n, 0, 1)*noise_sd
+y <- f_XW + rnorm(n, 0, 1) * noise_sd
 
 # Split data into test and train sets
 test_set_pct <- 0.2
-n_test <- round(test_set_pct*n)
+n_test <- round(test_set_pct * n)
 n_train <- n - n_test
 test_inds <- sort(sample(1:n, n_test, replace = FALSE))
 train_inds <- (1:n)[!((1:n) %in% test_inds)]
-X_test <- as.data.frame(X[test_inds,])
-X_train <- as.data.frame(X[train_inds,])
+X_test <- as.data.frame(X[test_inds, ])
+X_train <- as.data.frame(X[train_inds, ])
 y_test <- y[test_inds]
 y_train <- y[train_inds]
 
 # Run the GFR algorithm
-xbart_params <- list(sample_sigma_global = T, 
-                     num_trees_mean = num_trees, alpha_mean = 0.99, 
-                     beta_mean = 1, max_depth_mean = -1, 
-                     min_samples_leaf_mean = 1, sample_sigma_leaf = F, 
-                     sigma_leaf_init = 1/num_trees)
+xbart_params <- list(
+  sample_sigma_global = T,
+  num_trees_mean = num_trees,
+  alpha_mean = 0.99,
+  beta_mean = 1,
+  max_depth_mean = -1,
+  min_samples_leaf_mean = 1,
+  sample_sigma_leaf = F,
+  sigma_leaf_init = 1 / num_trees
+)
 xbart_model <- stochtree::bart(
-    X_train = X_train, y_train = y_train, X_test = X_test, 
-    num_gfr = num_gfr, num_burnin = 0, num_mcmc = 0, params = xbart_params
+  X_train = X_train,
+  y_train = y_train,
+  X_test = X_test,
+  num_gfr = num_gfr,
+  num_burnin = 0,
+  num_mcmc = 0,
+  params = xbart_params
 )
-plot(rowMeans(xbart_model$y_hat_test), y_test); abline(0,1)
+plot(rowMeans(xbart_model$y_hat_test), y_test)
+abline(0, 1)
 cat(sqrt(mean((rowMeans(xbart_model$y_hat_test) - y_test)^2)), "\n")
-cat(mean((apply(xbart_model$y_hat_test, 1, quantile, probs=0.05) <= y_test) & (apply(xbart_model$y_hat_test, 1, quantile, probs=0.95) >= y_test)), "\n")
+cat(
+  mean(
+    (apply(xbart_model$y_hat_test, 1, quantile, probs = 0.05) <= y_test) &
+      (apply(xbart_model$y_hat_test, 1, quantile, probs = 0.95) >= y_test)
+  ),
+  "\n"
+)
 xbart_model_string <- stochtree::saveBARTModelToJsonString(xbart_model)
 
 # Parallel setup
@@ -51,20 +71,32 @@ cl <- makeCluster(ncores)
 registerDoParallel(cl)
 
 # Run the parallel BART MCMC samplers
-bart_model_outputs <- foreach (i = 1:num_chains) %dopar% {
+bart_model_outputs <- foreach(i = 1:num_chains) %dopar%
+  {
     random_seed <- i
-    bart_params <- list(sample_sigma_global = T, sample_sigma_leaf = T, 
-                        num_trees_mean = num_trees, random_seed = random_seed, 
-                        alpha_mean = 0.999, beta_mean = 1)
+    bart_params <- list(
+      sample_sigma_global = T,
+      sample_sigma_leaf = T,
+      num_trees_mean = num_trees,
+      random_seed = random_seed,
+      alpha_mean = 0.999,
+      beta_mean = 1
+    )
     bart_model <- stochtree::bart(
-        X_train = X_train, y_train = y_train, X_test = X_test, 
-        num_gfr = 0, num_burnin = num_burnin, num_mcmc = num_mcmc, params = bart_params, 
-        previous_model_json = xbart_model_string, warmstart_sample_num = num_gfr - i + 1, 
+      X_train = X_train,
+      y_train = y_train,
+      X_test = X_test,
+      num_gfr = 0,
+      num_burnin = num_burnin,
+      num_mcmc = num_mcmc,
+      params = bart_params,
+      previous_model_json = xbart_model_string,
+      warmstart_sample_num = num_gfr - i + 1,
     )
     bart_model_string <- stochtree::saveBARTModelToJsonString(bart_model)
     y_hat_test <- bart_model$y_hat_test
-    list(model=bart_model_string, yhat=y_hat_test)
-}
+    list(model = bart_model_string, yhat = y_hat_test)
+  }
 
 # Close the cluster connection
 stopCluster(cl)
@@ -73,43 +105,89 @@ stopCluster(cl)
 bart_model_strings <- list()
 bart_model_yhats <- matrix(NA, nrow = length(y_test), ncol = num_chains)
 for (i in 1:length(bart_model_outputs)) {
-    bart_model_strings[[i]] <- bart_model_outputs[[i]]$model
-    bart_model_yhats[,i] <- rowMeans(bart_model_outputs[[i]]$yhat)
+  bart_model_strings[[i]] <- bart_model_outputs[[i]]$model
+  bart_model_yhats[, i] <- rowMeans(bart_model_outputs[[i]]$yhat)
 }
 combined_bart <- createBARTModelFromCombinedJsonString(bart_model_strings)
 
 # Inspect the results
-yhat_combined <- predict(combined_bart, X_test)$y_hat
-par(mfrow = c(1,2))
+yhat_combined <- predict(combined_bart, X = X_test)$y_hat
+par(mfrow = c(1, 2))
 for (i in 1:num_chains) {
-    offset <- (i-1)*num_mcmc
-    inds_start <- offset + 1
-    inds_end <- offset + num_mcmc
-    plot(rowMeans(yhat_combined[,inds_start:inds_end]), bart_model_yhats[,i],
-         xlab = "deserialized", ylab = "original", 
-         main = paste0("Chain ", i, "\nPredictions"))
-    abline(0,1,col="red",lty=3,lwd=3)
+  offset <- (i - 1) * num_mcmc
+  inds_start <- offset + 1
+  inds_end <- offset + num_mcmc
+  plot(
+    rowMeans(yhat_combined[, inds_start:inds_end]),
+    bart_model_yhats[, i],
+    xlab = "deserialized",
+    ylab = "original",
+    main = paste0("Chain ", i, "\nPredictions")
+  )
+  abline(0, 1, col = "red", lty = 3, lwd = 3)
 }
 for (i in 1:num_chains) {
-    offset <- (i-1)*num_mcmc
-    inds_start <- offset + 1
-    inds_end <- offset + num_mcmc
-    plot(rowMeans(yhat_combined[,inds_start:inds_end]), y_test,
-         xlab = "predicted", ylab = "actual", 
-         main = paste0("Chain ", i, "\nPredictions"))
-    abline(0,1,col="red",lty=3,lwd=3)
-    cat(sqrt(mean((rowMeans(yhat_combined[,inds_start:inds_end]) - y_test)^2)), "\n")
-    cat(mean((apply(yhat_combined[,inds_start:inds_end], 1, quantile, probs=0.05) <= y_test) & (apply(yhat_combined[,inds_start:inds_end], 1, quantile, probs=0.95) >= y_test)), "\n")
+  offset <- (i - 1) * num_mcmc
+  inds_start <- offset + 1
+  inds_end <- offset + num_mcmc
+  plot(
+    rowMeans(yhat_combined[, inds_start:inds_end]),
+    y_test,
+    xlab = "predicted",
+    ylab = "actual",
+    main = paste0("Chain ", i, "\nPredictions")
+  )
+  abline(0, 1, col = "red", lty = 3, lwd = 3)
+  cat(
+    sqrt(mean((rowMeans(yhat_combined[, inds_start:inds_end]) - y_test)^2)),
+    "\n"
+  )
+  cat(
+    mean(
+      (apply(yhat_combined[, inds_start:inds_end], 1, quantile, probs = 0.05) <=
+        y_test) &
+        (apply(
+          yhat_combined[, inds_start:inds_end],
+          1,
+          quantile,
+          probs = 0.95
+        ) >=
+          y_test)
+    ),
+    "\n"
+  )
 }
-par(mfrow = c(1,1))
+par(mfrow = c(1, 1))
 
 # Compare to a single chain of MCMC samples initialized at root
-bart_params <- list(sample_sigma_global = T, sample_sigma_leaf = T, 
-                    num_trees_mean = num_trees, alpha_mean = 0.95, beta_mean = 2)
+bart_params <- list(
+  sample_sigma_global = T,
+  sample_sigma_leaf = T,
+  num_trees_mean = num_trees,
+  alpha_mean = 0.95,
+  beta_mean = 2
+)
 bart_model <- stochtree::bart(
-    X_train = X_train, y_train = y_train, X_test = X_test, 
-    num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, params = bart_params
+  X_train = X_train,
+  y_train = y_train,
+  X_test = X_test,
+  num_gfr = 0,
+  num_burnin = 0,
+  num_mcmc = num_mcmc,
+  params = bart_params
+)
+plot(
+  rowMeans(bart_model$y_hat_test),
+  y_test,
+  xlab = "predicted",
+  ylab = "actual"
 )
-plot(rowMeans(bart_model$y_hat_test), y_test, xlab = "predicted", ylab = "actual"); abline(0,1)
+abline(0, 1)
 cat(sqrt(mean((rowMeans(bart_model$y_hat_test) - y_test)^2)), "\n")
-cat(mean((apply(bart_model$y_hat_test, 1, quantile, probs=0.05) <= y_test) & (apply(bart_model$y_hat_test, 1, quantile, probs=0.95) >= y_test)), "\n")
+cat(
+  mean(
+    (apply(bart_model$y_hat_test, 1, quantile, probs = 0.05) <= y_test) &
+      (apply(bart_model$y_hat_test, 1, quantile, probs = 0.95) >= y_test)
+  ),
+  "\n"
+)