Updated R BART use of user-provided rfx parameters

Author: andrewherren

R/bart.R

@@ -706,11 +706,7 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
         if (is.null(rfx_group_parameter_prior_mean)) {
             xi_init <- matrix(rep(alpha_init, num_rfx_groups),num_rfx_components,num_rfx_groups)
         } else {
-            xi_init <- expand_dims_1d(rfx_group_parameter_prior_mean, num_rfx_components)
-            # If it's a vector, expand to matrix
-            if (is.vector(xi_init)) {
-                xi_init <- matrix(rep(xi_init, num_rfx_groups), num_rfx_components, num_rfx_groups)
-            }
+            xi_init <- expand_dims_2d(rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups)
         }
 
         if (is.null(rfx_working_parameter_prior_cov)) {

test/R/testthat/test-bart.R

@@ -336,3 +336,98 @@ test_that("BART Predictions", {
     expect_equal(train_preds_mean_cached, train_preds_mean_recomputed)
     expect_equal(train_preds_variance_cached, train_preds_variance_recomputed)
 })
+
+test_that("Random Effects BART", {
+    skip_on_cran()
+    
+    # Generate simulated data
+    n <- 100
+    p <- 5
+    p_w <- 2
+    X <- matrix(runif(n*p), ncol = p)
+    W <- matrix(runif(n*p_w), ncol = p_w)
+    f_XW <- (
+        ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5*W[,1]) + 
+        ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5*W[,1]) + 
+        ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5*W[,1]) + 
+        ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5*W[,1])
+    )
+    rfx_group_ids <- sample(1:2, size = n, replace = TRUE)
+    rfx_basis <- cbind(rep(1, n), runif(n))
+    num_rfx_components <- ncol(rfx_basis)
+    num_rfx_groups <- length(unique(rfx_group_ids))
+    rfx_coefs <- matrix(c(-5, 5, 1, -1), ncol = 2, byrow = T)
+    rfx_term <- rowSums(rfx_coefs[rfx_group_ids,] * rfx_basis)
+    noise_sd <- 1
+    y <- f_XW + rfx_term + 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,]
+    W_test <- W[test_inds,]
+    W_train <- W[train_inds,]
+    rfx_group_ids_test <- rfx_group_ids[test_inds]
+    rfx_group_ids_train <- rfx_group_ids[train_inds]
+    rfx_basis_test <- rfx_basis[test_inds,]
+    rfx_basis_train <- rfx_basis[train_inds,]
+    y_test <- y[test_inds]
+    y_train <- y[train_inds]
+    
+    # Specify no rfx parameters
+    general_param_list <- list()
+    mean_forest_param_list <- list(sample_sigma2_leaf = FALSE)
+    expect_no_error(
+        bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                           leaf_basis_train = W_train, leaf_basis_test = W_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, 
+                           general_params = general_param_list, 
+                           mean_forest_params = mean_forest_param_list)
+    )
+    
+    # Specify all rfx parameters as scalars
+    general_param_list <- list(rfx_working_parameter_prior_mean = 1.,
+                               rfx_group_parameter_prior_mean = 1.,
+                               rfx_working_parameter_prior_cov = 1.,
+                               rfx_group_parameter_prior_cov = 1.,
+                               rfx_variance_prior_shape = 1,
+                               rfx_variance_prior_scale = 1)
+    mean_forest_param_list <- list(sample_sigma2_leaf = FALSE)
+    expect_no_error(
+        bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                           leaf_basis_train = W_train, leaf_basis_test = W_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, 
+                           general_params = general_param_list, 
+                           mean_forest_params = mean_forest_param_list)
+    )
+    
+    # Specify all relevant rfx parameters as vectors
+    general_param_list <- list(rfx_working_parameter_prior_mean = c(1.,1.),
+                               rfx_group_parameter_prior_mean = c(1.,1.),
+                               rfx_working_parameter_prior_cov = diag(1.,2),
+                               rfx_group_parameter_prior_cov = diag(1.,2),
+                               rfx_variance_prior_shape = 1,
+                               rfx_variance_prior_scale = 1)
+    mean_forest_param_list <- list(sample_sigma2_leaf = FALSE)
+    expect_no_error(
+        bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                           leaf_basis_train = W_train, leaf_basis_test = W_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, 
+                           general_params = general_param_list, 
+                           mean_forest_params = mean_forest_param_list)
+    )
+})