Updated multi-chain in R to initialize different chains from different forests if a previous model json was provided

Author: andrewherren

R/bart.R

@@ -28,7 +28,7 @@
 #' @param num_burnin Number of "burn-in" iterations of the MCMC sampler. Default: 0.
 #' @param num_mcmc Number of "retained" iterations of the MCMC sampler. Default: 100.
 #' @param previous_model_json (Optional) JSON string containing a previous BART model. This can be used to "continue" a sampler interactively after inspecting the samples or to run parallel chains "warm-started" from existing forest samples. Default: `NULL`.
-#' @param previous_model_warmstart_sample_num (Optional) Sample number from `previous_model_json` that will be used to warmstart this BART sampler. One-indexed (so that the first sample is used for warm-start by setting `previous_model_warmstart_sample_num = 1`). Default: `NULL`.
+#' @param previous_model_warmstart_sample_num (Optional) Sample number from `previous_model_json` that will be used to warmstart this BART sampler. One-indexed (so that the first sample is used for warm-start by setting `previous_model_warmstart_sample_num = 1`). Default: `NULL`.  If `num_chains` in the `general_params` list is > 1, then each successive chain will be initialized from a different sample, counting backwards from `previous_model_warmstart_sample_num`. That is, if `previous_model_warmstart_sample_num = 10` and `num_chains = 4`, then chain 1 will be initialized from sample 10, chain 2 from sample 9, chain 3 from sample 8, and chain 4 from sample 7. If `previous_model_json` is provided but `previous_model_warmstart_sample_num` is NULL, the last sample in the previous model will be used to initialize the first chain, counting backwards as noted before. If more chains are requested than there are samples in `previous_model_json`, a warning will be raised and only the last sample will be used.
 #' @param general_params (Optional) A list of general (non-forest-specific) model parameters, each of which has a default value processed internally, so this argument list is optional.
 #'
 #'   - `cutpoint_grid_size` Maximum size of the "grid" of potential cutpoints to consider in the GFR algorithm. Default: `100`.
@@ -293,10 +293,36 @@ bart <- function(
 
   # Check if previous model JSON is provided and parse it if so
   has_prev_model <- !is.null(previous_model_json)
+  has_prev_model_index <- !is.null(previous_model_warmstart_sample_num)
   if (has_prev_model) {
     previous_bart_model <- createBARTModelFromJsonString(
       previous_model_json
     )
+    prev_num_samples <- previous_bart_model$model_params$num_samples
+    if (!has_prev_model_index) {
+      previous_model_warmstart_sample_num <- prev_num_samples
+      warning(
+        "`previous_model_warmstart_sample_num` was not provided alongside `previous_model_json`, so it will be set to the number of samples available in `previous_model_json`"
+      )
+    } else {
+      if (previous_model_warmstart_sample_num < 1) {
+        stop(
+          "`previous_model_warmstart_sample_num` must be a positive integer"
+        )
+      }
+      if (previous_model_warmstart_sample_num > prev_num_samples) {
+        stop(
+          "`previous_model_warmstart_sample_num` exceeds the number of samples in `previous_model_json`"
+        )
+      }
+    }
+    previous_model_decrement <- T
+    if (num_chains > previous_model_warmstart_sample_num) {
+      warning(
+        "The number of chains being sampled exceeds the number of previous model samples available from the requested position in `previous_model_json`. All chains will be initialized from the same sample."
+      )
+      previous_model_decrement <- F
+    }
     previous_y_bar <- previous_bart_model$model_params$outcome_mean
     previous_y_scale <- previous_bart_model$model_params$outcome_scale
     if (previous_bart_model$model_params$include_mean_forest) {
@@ -1375,11 +1401,16 @@ bart <- function(
           )
         }
       } else if (has_prev_model) {
+        warmstart_index <- ifelse(
+          previous_model_decrement,
+          previous_model_warmstart_sample_num - chain_num + 1,
+          previous_model_warmstart_sample_num
+        )
         if (include_mean_forest) {
           resetActiveForest(
             active_forest_mean,
             previous_forest_samples_mean,
-            previous_model_warmstart_sample_num - 1
+            warmstart_index - 1
           )
           resetForestModel(
             forest_model_mean,
@@ -1393,7 +1424,7 @@ bart <- function(
               (!is.null(previous_leaf_var_samples))
           ) {
             leaf_scale_double <- previous_leaf_var_samples[
-              previous_model_warmstart_sample_num
+              warmstart_index
             ]
             current_leaf_scale <- as.matrix(leaf_scale_double)
             forest_model_config_mean$update_leaf_model_scale(
@@ -1405,7 +1436,7 @@ bart <- function(
           resetActiveForest(
             active_forest_variance,
             previous_forest_samples_variance,
-            previous_model_warmstart_sample_num - 1
+            warmstart_index - 1
           )
           resetForestModel(
             forest_model_variance,
@@ -1439,7 +1470,7 @@ bart <- function(
             resetRandomEffectsModel(
               rfx_model,
               previous_rfx_samples,
-              previous_model_warmstart_sample_num - 1,
+              warmstart_index - 1,
               sigma_alpha_init
             )
             resetRandomEffectsTracker(
@@ -1454,7 +1485,7 @@ bart <- function(
         if (sample_sigma2_global) {
           if (!is.null(previous_global_var_samples)) {
             current_sigma2 <- previous_global_var_samples[
-              previous_model_warmstart_sample_num
+              warmstart_index
             ]
             global_model_config$update_global_error_variance(
               current_sigma2

R/bcf.R

@@ -25,7 +25,7 @@
 #' @param num_burnin Number of "burn-in" iterations of the MCMC sampler. Default: 0.
 #' @param num_mcmc Number of "retained" iterations of the MCMC sampler. Default: 100.
 #' @param previous_model_json (Optional) JSON string containing a previous BCF model. This can be used to "continue" a sampler interactively after inspecting the samples or to run parallel chains "warm-started" from existing forest samples. Default: `NULL`.
-#' @param previous_model_warmstart_sample_num (Optional) Sample number from `previous_model_json` that will be used to warmstart this BCF sampler. One-indexed (so that the first sample is used for warm-start by setting `previous_model_warmstart_sample_num = 1`). Default: `NULL`.
+#' @param previous_model_warmstart_sample_num (Optional) Sample number from `previous_model_json` that will be used to warmstart this BCF sampler. One-indexed (so that the first sample is used for warm-start by setting `previous_model_warmstart_sample_num = 1`). Default: `NULL`. If `num_chains` in the `general_params` list is > 1, then each successive chain will be initialized from a different sample, counting backwards from `previous_model_warmstart_sample_num`. That is, if `previous_model_warmstart_sample_num = 10` and `num_chains = 4`, then chain 1 will be initialized from sample 10, chain 2 from sample 9, chain 3 from sample 8, and chain 4 from sample 7. If `previous_model_json` is provided but `previous_model_warmstart_sample_num` is NULL, the last sample in the previous model will be used to initialize the first chain, counting backwards as noted before. If more chains are requested than there are samples in `previous_model_json`, a warning will be raised and only the last sample will be used.
 #' @param general_params (Optional) A list of general (non-forest-specific) model parameters, each of which has a default value processed internally, so this argument list is optional.
 #'
 #'   - `cutpoint_grid_size` Maximum size of the "grid" of potential cutpoints to consider in the GFR algorithm. Default: `100`.
@@ -397,8 +397,34 @@ bcf <- function(
 
   # Check if previous model JSON is provided and parse it if so
   has_prev_model <- !is.null(previous_model_json)
+  has_prev_model_index <- !is.null(previous_model_warmstart_sample_num)
   if (has_prev_model) {
     previous_bcf_model <- createBCFModelFromJsonString(previous_model_json)
+    prev_num_samples <- previous_bcf_model$model_params$num_samples
+    if (!has_prev_model_index) {
+      previous_model_warmstart_sample_num <- prev_num_samples
+      warning(
+        "`previous_model_warmstart_sample_num` was not provided alongside `previous_model_json`, so it will be set to the number of samples available in `previous_model_json`"
+      )
+    } else {
+      if (previous_model_warmstart_sample_num < 1) {
+        stop(
+          "`previous_model_warmstart_sample_num` must be a positive integer"
+        )
+      }
+      if (previous_model_warmstart_sample_num > prev_num_samples) {
+        stop(
+          "`previous_model_warmstart_sample_num` exceeds the number of samples in `previous_model_json`"
+        )
+      }
+    }
+    previous_model_decrement <- T
+    if (num_chains > previous_model_warmstart_sample_num) {
+      warning(
+        "The number of chains being sampled exceeds the number of previous model samples available from the requested position in `previous_model_json`. All chains will be initialized from the same sample."
+      )
+      previous_model_decrement <- F
+    }
     previous_y_bar <- previous_bcf_model$model_params$outcome_mean
     previous_y_scale <- previous_bcf_model$model_params$outcome_scale
     previous_forest_samples_mu <- previous_bcf_model$forests_mu
@@ -1910,10 +1936,15 @@ bcf <- function(
           )
         }
       } else if (has_prev_model) {
+        warmstart_index <- ifelse(
+          previous_model_decrement,
+          previous_model_warmstart_sample_num - chain_num + 1,
+          previous_model_warmstart_sample_num
+        )
         resetActiveForest(
           active_forest_mu,
           previous_forest_samples_mu,
-          previous_model_warmstart_sample_num - 1
+          warmstart_index - 1
         )
         resetForestModel(
           forest_model_mu,
@@ -1925,7 +1956,7 @@ bcf <- function(
         resetActiveForest(
           active_forest_tau,
           previous_forest_samples_tau,
-          previous_model_warmstart_sample_num - 1
+          warmstart_index - 1
         )
         resetForestModel(
           forest_model_tau,
@@ -1938,7 +1969,7 @@ bcf <- function(
           resetActiveForest(
             active_forest_variance,
             previous_forest_samples_variance,
-            previous_model_warmstart_sample_num - 1
+            warmstart_index - 1
           )
           resetForestModel(
             forest_model_variance,
@@ -1953,7 +1984,7 @@ bcf <- function(
             (!is.null(previous_leaf_var_mu_samples))
         ) {
           leaf_scale_mu_double <- previous_leaf_var_mu_samples[
-            previous_model_warmstart_sample_num
+            warmstart_index
           ]
           current_leaf_scale_mu <- as.matrix(leaf_scale_mu_double)
           forest_model_config_mu$update_leaf_model_scale(
@@ -1965,7 +1996,7 @@ bcf <- function(
             (!is.null(previous_leaf_var_tau_samples))
         ) {
           leaf_scale_tau_double <- previous_leaf_var_tau_samples[
-            previous_model_warmstart_sample_num
+            warmstart_index
           ]
           current_leaf_scale_tau <- as.matrix(leaf_scale_tau_double)
           forest_model_config_tau$update_leaf_model_scale(
@@ -1975,12 +2006,12 @@ bcf <- function(
         if (adaptive_coding) {
           if (!is.null(previous_b_1_samples)) {
             current_b_1 <- previous_b_1_samples[
-              previous_model_warmstart_sample_num
+              warmstart_index
             ]
           }
           if (!is.null(previous_b_0_samples)) {
             current_b_0 <- previous_b_0_samples[
-              previous_model_warmstart_sample_num
+              warmstart_index
             ]
           }
           tau_basis_train <- (1 - Z_train) *
@@ -2023,7 +2054,7 @@ bcf <- function(
             resetRandomEffectsModel(
               rfx_model,
               previous_rfx_samples,
-              previous_model_warmstart_sample_num - 1,
+              warmstart_index - 1,
               sigma_alpha_init
             )
             resetRandomEffectsTracker(
@@ -2038,7 +2069,7 @@ bcf <- function(
         if (sample_sigma2_global) {
           if (!is.null(previous_global_var_samples)) {
             current_sigma2 <- previous_global_var_samples[
-              previous_model_warmstart_sample_num
+              warmstart_index
             ]
           }
           global_model_config$update_global_error_variance(