Fixed minor BCF bug

Author: andrewherren

R/bcf.R

@@ -286,8 +286,8 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
             previous_forest_samples_variance <- previous_bcf_model$forests_variance
         } else previous_forest_samples_variance <- NULL
         if (previous_bcf_model$model_params$sample_sigma_global) {
-            previous_global_var_samples <- previous_bcf_model$sigma2_samples*(
-                previous_var_scale / (previous_y_scale*previous_y_scale)
+            previous_global_var_samples <- previous_bcf_model$sigma2_samples / (
+                previous_y_scale*previous_y_scale
             )
         } else previous_global_var_samples <- NULL
         if (previous_bcf_model$model_params$sample_sigma_leaf_mu) {
@@ -313,7 +313,6 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
     } else {
         previous_y_bar <- NULL
         previous_y_scale <- NULL
-        previous_var_scale <- NULL
         previous_global_var_samples <- NULL
         previous_leaf_var_mu_samples <- NULL
         previous_leaf_var_tau_samples <- NULL

test/R/testthat/test-bcf.R

@@ -81,7 +81,7 @@ test_that("MCMC BCF", {
     )
 })
 
-test_that("GFR BART", {
+test_that("GFR BCF", {
     skip_on_cran()
 
     # Generate simulated data
@@ -90,21 +90,21 @@ test_that("GFR BART", {
     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)
+        ((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)
+        ((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)
+        ((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
@@ -181,3 +181,151 @@ test_that("GFR BART", {
                          num_mcmc = 10, general_params = general_param_list)
     )
 })
+
+test_that("Warmstart BCF", {
+    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[test_inds]
+    pi_train <- pi[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_param_list <- list(num_chains = 1, keep_every = 1)
+    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 = 0, general_params = general_param_list)
+    
+    # Save to JSON string
+    bcf_model_json_string <- saveBCFModelToJsonString(bcf_model)
+    
+    # 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 = 1, 
+                         general_params = general_param_list)
+    )
+    
+    # Generate simulated data with random effects
+    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)
+    rfx_group_ids <- sample(1:2, size = n, replace = T)
+    rfx_basis <- rep(1, n)
+    rfx_coefs <- c(-5, 5)
+    rfx_term <- rfx_coefs[rfx_group_ids] * rfx_basis
+    noise_sd <- 1
+    y <- mu_X + tau_X*Z + 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,]
+    Z_test <- Z[test_inds]
+    Z_train <- Z[train_inds]
+    pi_test <- pi[test_inds]
+    pi_train <- pi[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]
+    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]
+    
+    # Run a BCF model with only GFR
+    general_param_list <- list(num_chains = 1, keep_every = 1)
+    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,  
+                     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, 
+                     propensity_test = pi_test, num_gfr = 10, num_burnin = 0, 
+                     num_mcmc = 0, general_params = general_param_list)
+    
+    # Save to JSON string
+    bcf_model_json_string <- saveBCFModelToJsonString(bcf_model)
+    
+    # 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,  
+                         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, 
+                         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 = 1, 
+                         general_params = general_param_list)
+    )
+})

test/R/testthat/test-serialization.R

@@ -7,9 +7,9 @@ test_that("BART Serialization", {
     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)
+        ((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)