Implemented and tested higher-level "collapse" functionality for forest containers in R

Author: andrewherren

R/forest.R

@@ -23,7 +23,42 @@ ForestSamples <- R6::R6Class(
         }, 
 
         #' @description
-        #' Collapse specified forests into a single forest
+        #' Collapse forests in this container by a pre-specified batch size. 
+        #' For example, if we have a container of twenty 10-tree forests, and we 
+        #' specify a `batch_size` of 5, then this method will yield four 50-tree 
+        #' forests. "Excess" forests remaining after the size of a forest container 
+        #' is divided by `batch_size` will be pruned from the beginning of the 
+        #' container (i.e. earlier sampled forests will be deleted). This method 
+        #' has no effect if `batch_size` is larger than the number of forests 
+        #' in a container.
+        #' @param batch_size Number of forests to be collapsed into a single forest
+        collapse = function(batch_size) {
+            container_size <- self$num_samples()
+            if ((batch_size <= container_size) && (batch_size > 1)) {
+                reverse_container_inds <- seq(container_size, 1, -1)
+                num_clean_batches <- container_size %/% batch_size
+                batch_inds <- (reverse_container_inds - (container_size - (container_size %/% num_clean_batches) * num_clean_batches) - 1) %/% batch_size
+                for (batch_ind in unique(batch_inds[batch_inds >= 0])) {
+                    merge_forest_inds <- sort(reverse_container_inds[batch_inds == batch_ind] - 1)
+                    num_merge_forests <- length(merge_forest_inds)
+                    self$combine_forests(merge_forest_inds)
+                    for (i in num_merge_forests:2) {
+                        self$delete_sample(merge_forest_inds[i])
+                    }
+                    forest_scale_factor <- 1.0 / num_merge_forests
+                    self$multiply_forest(merge_forest_inds[1], forest_scale_factor)
+                }
+                if (min(batch_inds) < 0) {
+                    delete_forest_inds <- sort(reverse_container_inds[batch_inds < 0] - 1)
+                    for (i in length(delete_forest_inds):1) {
+                        self$delete_sample(delete_forest_inds[i])
+                    }
+                }
+            }
+        }, 
+        
+        #' @description
+        #' Merge specified forests into a single forest
         #' @param forest_inds Indices of forests to be combined (0-indexed)
         combine_forests = function(forest_inds) {
             stopifnot(max(forest_inds) < self$num_samples())

man/ForestSamples.Rd

@@ -108,3 +108,144 @@ test_that("Univariate constant forest container", {
     # Assertion
     expect_equal(pred, pred_expected_new)
 })
+
+test_that("Collapse forests", {
+    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]
+    
+    # Create forest dataset
+    forest_dataset_test <- createForestDataset(covariates = X_test)
+    
+    # Run BART for 50 iterations
+    num_mcmc <- 50
+    general_param_list <- list(num_chains = 1, keep_every = 1)
+    bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                       num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, 
+                       general_params = general_param_list)
+    
+    # Extract the mean forest container
+    mean_forest_container <- bart_model$mean_forests
+    
+    # Predict from the original container
+    pred_orig <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Collapse the container in batches of 5
+    batch_size <- 5
+    mean_forest_container$collapse(batch_size)
+    
+    # Predict from the modified container
+    pred_new <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Check that corresponding (sums of) predictions match
+    batch_inds <- (seq(1,num_mcmc,1) - (num_mcmc - (num_mcmc %/% (num_mcmc %/% batch_size)) * (num_mcmc %/% batch_size)) - 1) %/% batch_size + 1
+    pred_orig_collapsed <- matrix(NA, nrow = nrow(pred_orig), ncol = max(batch_inds))
+    for (i in 1:max(batch_inds)) {
+        pred_orig_collapsed[,i] <- rowSums(pred_orig[,batch_inds == i]) / sum(batch_inds == i)
+    }
+    
+    # Assertion
+    expect_equal(pred_orig_collapsed, pred_new)
+    
+    # Now run BART for 52 iterations
+    num_mcmc <- 52
+    general_param_list <- list(num_chains = 1, keep_every = 1)
+    bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                       num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, 
+                       general_params = general_param_list)
+    
+    # Extract the mean forest container
+    mean_forest_container <- bart_model$mean_forests
+    
+    # Predict from the original container
+    pred_orig <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Collapse the container in batches of 5
+    batch_size <- 5
+    mean_forest_container$collapse(batch_size)
+    
+    # Predict from the modified container
+    pred_new <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Check that corresponding (sums of) predictions match
+    batch_inds <- (seq(1,num_mcmc,1) - (num_mcmc - (num_mcmc %/% (num_mcmc %/% batch_size)) * (num_mcmc %/% batch_size)) - 1) %/% batch_size + 1
+    pred_orig_collapsed <- matrix(NA, nrow = nrow(pred_orig), ncol = max(batch_inds))
+    for (i in 1:max(batch_inds)) {
+        pred_orig_collapsed[,i] <- rowSums(pred_orig[,batch_inds == i]) / sum(batch_inds == i)
+    }
+    
+    # Assertion
+    expect_equal(pred_orig_collapsed, pred_new)
+    
+    # Now run BART for 5 iterations
+    num_mcmc <- 5
+    general_param_list <- list(num_chains = 1, keep_every = 1)
+    bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                       num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, 
+                       general_params = general_param_list)
+    
+    # Extract the mean forest container
+    mean_forest_container <- bart_model$mean_forests
+    
+    # Predict from the original container
+    pred_orig <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Collapse the container in batches of 5
+    batch_size <- 5
+    mean_forest_container$collapse(batch_size)
+    
+    # Predict from the modified container
+    pred_new <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Check that corresponding (sums of) predictions match
+    pred_orig_collapsed <- as.matrix(rowSums(pred_orig) / batch_size)
+
+    # Assertion
+    expect_equal(pred_orig_collapsed, pred_new)
+    
+    # Now run BART for 4 iterations
+    num_mcmc <- 4
+    general_param_list <- list(num_chains = 1, keep_every = 1)
+    bart_model <- bart(X_train = X_train, y_train = y_train, X_test = X_test, 
+                       num_gfr = 0, num_burnin = 0, num_mcmc = num_mcmc, 
+                       general_params = general_param_list)
+    
+    # Extract the mean forest container
+    mean_forest_container <- bart_model$mean_forests
+    
+    # Predict from the original container
+    pred_orig <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Collapse the container in batches of 5
+    batch_size <- 5
+    mean_forest_container$collapse(batch_size)
+    
+    # Predict from the modified container
+    pred_new <- mean_forest_container$predict(forest_dataset_test)
+    
+    # Check that corresponding (sums of) predictions match
+    pred_orig_collapsed <- pred_orig
+    
+    # Assertion
+    expect_equal(pred_orig_collapsed, pred_new)
+})