Added R BART and BCF unit tests for cached predictions

Author: andrewherren

test/R/testthat/test-bart.R

@@ -291,3 +291,48 @@ test_that("Warmstart BART", {
                            general_params = general_param_list)
     )
 })
+
+test_that("BART Predictions", {
+    skip_on_cran()
+    
+    # Generate simulated data
+    n <- 100
+    p <- 5
+    X <- matrix(runif(n*p), ncol = p)
+    f_XW <- (
+        ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + 
+            ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + 
+            ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + 
+            ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5)
+    )
+    noise_sd <- 1
+    y <- f_XW + rnorm(n, 0, noise_sd)
+    test_set_pct <- 0.2
+    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 <- X[test_inds,]
+    X_train <- X[train_inds,]
+    y_test <- y[test_inds]
+    y_train <- y[train_inds]
+    
+    # Run a BART model with only GFR
+    general_params <- list(num_chains = 1)
+    variance_forest_params <- list(num_trees = 50)
+    bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                       num_gfr = 10, num_burnin = 0, num_mcmc = 10, 
+                       general_params = general_params, 
+                       variance_forest_params = variance_forest_params)
+    
+    # Check that cached predictions agree with results of predict() function
+    train_preds <- predict(bart_model, X = X_train)
+    train_preds_mean_cached <- bart_model$y_hat_train
+    train_preds_mean_recomputed <- train_preds$mean_forest_predictions
+    train_preds_variance_cached <- bart_model$sigma2_x_hat_train
+    train_preds_variance_recomputed <- train_preds$variance_forest_predictions
+    
+    # Assertion
+    expect_equal(train_preds_mean_cached, train_preds_mean_recomputed)
+    expect_equal(train_preds_variance_cached, train_preds_variance_recomputed)
+})

test/R/testthat/test-bcf.R

@@ -426,4 +426,71 @@ test_that("Multivariate Treatment MCMC BCF", {
                          propensity_test = pi_test, num_gfr = 0, num_burnin = 10, 
                          num_mcmc = 10, general_params = general_param_list)
     )
-})
+})
+
+test_that("BCF Predictions", {
+    skip_on_cran()
+    
+    # Generate simulated data
+    n <- 100
+    p <- 5
+    X <- matrix(runif(n*p), ncol = p)
+    mu_X <- (
+        ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) + 
+            ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) + 
+            ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) + 
+            ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5)
+    )
+    pi_X <- (
+        ((0 <= X[,1]) & (0.25 > X[,1])) * (0.2) + 
+            ((0.25 <= X[,1]) & (0.5 > X[,1])) * (0.4) + 
+            ((0.5 <= X[,1]) & (0.75 > X[,1])) * (0.6) + 
+            ((0.75 <= X[,1]) & (1 > X[,1])) * (0.8)
+    )
+    tau_X <- (
+        ((0 <= X[,2]) & (0.25 > X[,2])) * (0.5) + 
+            ((0.25 <= X[,2]) & (0.5 > X[,2])) * (1.0) + 
+            ((0.5 <= X[,2]) & (0.75 > X[,2])) * (1.5) + 
+            ((0.75 <= X[,2]) & (1 > X[,2])) * (2.0)
+    )
+    Z <- rbinom(n, 1, pi_X)
+    noise_sd <- 1
+    y <- mu_X + tau_X*Z + rnorm(n, 0, noise_sd)
+    test_set_pct <- 0.2
+    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 <- X[test_inds,]
+    X_train <- X[train_inds,]
+    Z_test <- Z[test_inds]
+    Z_train <- Z[train_inds]
+    pi_test <- pi_X[test_inds]
+    pi_train <- pi_X[train_inds]
+    mu_test <- mu_X[test_inds]
+    mu_train <- mu_X[train_inds]
+    tau_test <- tau_X[test_inds]
+    tau_train <- tau_X[train_inds]
+    y_test <- y[test_inds]
+    y_train <- y[train_inds]
+    
+    # Run a BCF model with only GFR
+    general_params <- list(num_chains = 1, keep_every = 1)
+    variance_forest_params <- list(num_trees = 50)
+    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 = 10, num_burnin = 0, 
+                     num_mcmc = 10, general_params = general_params, 
+                     variance_forest_params = variance_forest_params)
+    
+    # Check that cached predictions agree with results of predict() function
+    train_preds <- predict(bcf_model, X = X_train, Z = Z_train, propensity = pi_train)
+    train_preds_mean_cached <- bcf_model$y_hat_train
+    train_preds_mean_recomputed <- train_preds$y_hat
+    train_preds_variance_cached <- bcf_model$sigma2_x_hat_train
+    train_preds_variance_recomputed <- train_preds$variance_forest_predictions
+    
+    # Assertion
+    expect_equal(train_preds_mean_cached, train_preds_mean_recomputed)
+    expect_equal(train_preds_variance_cached, train_preds_variance_recomputed)
+})