Enable multi-chain sampling (serial or parallel)

.gitignore

@@ -1,14 +1,16 @@
 ## System and data files
 *.pdf
 *.csv
-*.txt
 *.DS_Store
 lib/
 build/
 .vscode/
 xcode/
 *.json
 .vs/
+cpp_docs/doxyoutput/html
+cpp_docs/doxyoutput/xml
+cpp_docs/doxyoutput/latex
 
 ## R gitignore
 

NAMESPACE

@@ -9,11 +9,19 @@ export(bcf)
 export(calibrate_inverse_gamma_error_variance)
 export(computeForestKernels)
 export(computeForestLeafIndices)
+export(convertBARTModelToJson)
 export(convertBCFModelToJson)
+export(createBARTModelFromCombinedJson)
+export(createBARTModelFromCombinedJsonString)
+export(createBARTModelFromJson)
+export(createBARTModelFromJsonFile)
+export(createBARTModelFromJsonString)
 export(createBCFModelFromJson)
 export(createBCFModelFromJsonFile)
+export(createBCFModelFromJsonString)
 export(createCppJson)
 export(createCppJsonFile)
+export(createCppJsonString)
 export(createForestContainer)
 export(createForestCovariates)
 export(createForestCovariatesFromMetadata)
@@ -27,7 +35,11 @@ export(createRandomEffectsDataset)
 export(createRandomEffectsModel)
 export(createRandomEffectsTracker)
 export(getRandomEffectSamples)
+export(loadForestContainerCombinedJson)
+export(loadForestContainerCombinedJsonString)
 export(loadForestContainerJson)
+export(loadRandomEffectSamplesCombinedJson)
+export(loadRandomEffectSamplesCombinedJsonString)
 export(loadRandomEffectSamplesJson)
 export(loadScalarJson)
 export(loadVectorJson)
@@ -43,7 +55,10 @@ export(preprocessTrainDataFrame)
 export(preprocessTrainMatrix)
 export(sample_sigma2_one_iteration)
 export(sample_tau_one_iteration)
+export(saveBARTModelToJsonFile)
+export(saveBARTModelToJsonString)
 export(saveBCFModelToJsonFile)
+export(saveBCFModelToJsonString)
 importFrom(R6,R6Class)
 importFrom(stats,lm)
 importFrom(stats,model.matrix)

R/cpp11.R

@@ -104,6 +104,26 @@ rfx_group_ids_from_json_cpp <- function(json_ptr, rfx_label) {
   .Call(`_stochtree_rfx_group_ids_from_json_cpp`, json_ptr, rfx_label)
 }
 
+rfx_container_append_from_json_cpp <- function(rfx_container_ptr, json_ptr, rfx_label) {
+  invisible(.Call(`_stochtree_rfx_container_append_from_json_cpp`, rfx_container_ptr, json_ptr, rfx_label))
+}
+
+rfx_container_from_json_string_cpp <- function(json_string, rfx_label) {
+  .Call(`_stochtree_rfx_container_from_json_string_cpp`, json_string, rfx_label)
+}
+
+rfx_label_mapper_from_json_string_cpp <- function(json_string, rfx_label) {
+  .Call(`_stochtree_rfx_label_mapper_from_json_string_cpp`, json_string, rfx_label)
+}
+
+rfx_group_ids_from_json_string_cpp <- function(json_string, rfx_label) {
+  .Call(`_stochtree_rfx_group_ids_from_json_string_cpp`, json_string, rfx_label)
+}
+
+rfx_container_append_from_json_string_cpp <- function(rfx_container_ptr, json_string, rfx_label) {
+  invisible(.Call(`_stochtree_rfx_container_append_from_json_string_cpp`, rfx_container_ptr, json_string, rfx_label))
+}
+
 rfx_model_cpp <- function(num_components, num_groups) {
   .Call(`_stochtree_rfx_model_cpp`, num_components, num_groups)
 }
@@ -188,14 +208,26 @@ rfx_label_mapper_to_list_cpp <- function(label_mapper_ptr) {
   .Call(`_stochtree_rfx_label_mapper_to_list_cpp`, label_mapper_ptr)
 }
 
-forest_container_cpp <- function(num_trees, output_dimension, is_leaf_constant) {
-  .Call(`_stochtree_forest_container_cpp`, num_trees, output_dimension, is_leaf_constant)
+forest_container_cpp <- function(num_trees, output_dimension, is_leaf_constant, is_exponentiated) {
+  .Call(`_stochtree_forest_container_cpp`, num_trees, output_dimension, is_leaf_constant, is_exponentiated)
 }
 
 forest_container_from_json_cpp <- function(json_ptr, forest_label) {
   .Call(`_stochtree_forest_container_from_json_cpp`, json_ptr, forest_label)
 }
 
+forest_container_append_from_json_cpp <- function(forest_sample_ptr, json_ptr, forest_label) {
+  invisible(.Call(`_stochtree_forest_container_append_from_json_cpp`, forest_sample_ptr, json_ptr, forest_label))
+}
+
+forest_container_from_json_string_cpp <- function(json_string, forest_label) {
+  .Call(`_stochtree_forest_container_from_json_string_cpp`, json_string, forest_label)
+}
+
+forest_container_append_from_json_string_cpp <- function(forest_sample_ptr, json_string, forest_label) {
+  invisible(.Call(`_stochtree_forest_container_append_from_json_string_cpp`, forest_sample_ptr, json_string, forest_label))
+}
+
 num_samples_forest_container_cpp <- function(forest_samples) {
   .Call(`_stochtree_num_samples_forest_container_cpp`, forest_samples)
 }
@@ -284,6 +316,10 @@ set_leaf_vector_forest_container_cpp <- function(forest_samples, leaf_vector) {
   invisible(.Call(`_stochtree_set_leaf_vector_forest_container_cpp`, forest_samples, leaf_vector))
 }
 
+initialize_forest_model_cpp <- function(data, residual, forest_samples, tracker, init_values, leaf_model_int) {
+  invisible(.Call(`_stochtree_initialize_forest_model_cpp`, data, residual, forest_samples, tracker, init_values, leaf_model_int))
+}
+
 adjust_residual_forest_container_cpp <- function(data, residual, forest_samples, tracker, requires_basis, forest_num, add) {
   invisible(.Call(`_stochtree_adjust_residual_forest_container_cpp`, data, residual, forest_samples, tracker, requires_basis, forest_num, add))
 }
@@ -332,16 +368,16 @@ forest_kernel_compute_kernel_train_test_cpp <- function(forest_kernel, covariate
   .Call(`_stochtree_forest_kernel_compute_kernel_train_test_cpp`, forest_kernel, covariates_train, covariates_test, forest_container, forest_num)
 }
 
-sample_gfr_one_iteration_cpp <- function(data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized) {
-  invisible(.Call(`_stochtree_sample_gfr_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized))
+sample_gfr_one_iteration_cpp <- function(data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, a_forest, b_forest, global_variance, leaf_model_int, pre_initialized) {
+  invisible(.Call(`_stochtree_sample_gfr_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, a_forest, b_forest, global_variance, leaf_model_int, pre_initialized))
 }
 
-sample_mcmc_one_iteration_cpp <- function(data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized) {
-  invisible(.Call(`_stochtree_sample_mcmc_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized))
+sample_mcmc_one_iteration_cpp <- function(data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, a_forest, b_forest, global_variance, leaf_model_int, pre_initialized) {
+  invisible(.Call(`_stochtree_sample_mcmc_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, a_forest, b_forest, global_variance, leaf_model_int, pre_initialized))
 }
 
-sample_sigma2_one_iteration_cpp <- function(residual, rng, a, b) {
-  .Call(`_stochtree_sample_sigma2_one_iteration_cpp`, residual, rng, a, b)
+sample_sigma2_one_iteration_cpp <- function(residual, dataset, rng, a, b) {
+  .Call(`_stochtree_sample_sigma2_one_iteration_cpp`, residual, dataset, rng, a, b)
 }
 
 sample_tau_one_iteration_cpp <- function(forest_samples, rng, a, b, sample_num) {
@@ -472,10 +508,18 @@ json_add_rfx_groupids_cpp <- function(json_ptr, groupids) {
   .Call(`_stochtree_json_add_rfx_groupids_cpp`, json_ptr, groupids)
 }
 
-json_save_cpp <- function(json_ptr, filename) {
-  invisible(.Call(`_stochtree_json_save_cpp`, json_ptr, filename))
+get_json_string_cpp <- function(json_ptr) {
+  .Call(`_stochtree_get_json_string_cpp`, json_ptr)
+}
+
+json_save_file_cpp <- function(json_ptr, filename) {
+  invisible(.Call(`_stochtree_json_save_file_cpp`, json_ptr, filename))
+}
+
+json_load_file_cpp <- function(json_ptr, filename) {
+  invisible(.Call(`_stochtree_json_load_file_cpp`, json_ptr, filename))
 }
 
-json_load_cpp <- function(json_ptr, filename) {
-  invisible(.Call(`_stochtree_json_load_cpp`, json_ptr, filename))
+json_load_string_cpp <- function(json_ptr, json_string) {
+  invisible(.Call(`_stochtree_json_load_string_cpp`, json_ptr, json_string))
 }

R/forest.R

@@ -16,20 +16,48 @@ ForestSamples <- R6::R6Class(
         #' @param num_trees Number of trees
         #' @param output_dimension Dimensionality of the outcome model
         #' @param is_leaf_constant Whether leaf is constant
+        #' @param is_exponentiated Whether forest predictions should be exponentiated before being returned
         #' @return A new `ForestContainer` object.
-        initialize = function(num_trees, output_dimension=1, is_leaf_constant=F) {
-            self$forest_container_ptr <- forest_container_cpp(num_trees, output_dimension, is_leaf_constant)
+        initialize = function(num_trees, output_dimension=1, is_leaf_constant=F, is_exponentiated=F) {
+            self$forest_container_ptr <- forest_container_cpp(num_trees, output_dimension, is_leaf_constant, is_exponentiated)
         }, 
 
         #' @description
-        #' Create a new ForestContainer object from a json object
+        #' Create a new `ForestContainer` object from a json object
         #' @param json_object Object of class `CppJson`
         #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy
         #' @return A new `ForestContainer` object.
         load_from_json = function(json_object, json_forest_label) {
             self$forest_container_ptr <- forest_container_from_json_cpp(json_object$json_ptr, json_forest_label)
         }, 
 
+        #' @description
+        #' Append to a `ForestContainer` object from a json object
+        #' @param json_object Object of class `CppJson`
+        #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy
+        #' @return NULL
+        append_from_json = function(json_object, json_forest_label) {
+            forest_container_append_from_json_cpp(self$forest_container_ptr, json_object$json_ptr, json_forest_label)
+        }, 
+
+        #' @description
+        #' Create a new `ForestContainer` object from a json object
+        #' @param json_string JSON string which parses into object of class `CppJson`
+        #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy
+        #' @return A new `ForestContainer` object.
+        load_from_json_string = function(json_string, json_forest_label) {
+            self$forest_container_ptr <- forest_container_from_json_string_cpp(json_string, json_forest_label)
+        }, 
+
+        #' @description
+        #' Append to a `ForestContainer` object from a json object
+        #' @param json_string JSON string which parses into object of class `CppJson`
+        #' @param json_forest_label Label referring to a particular forest (i.e. "forest_0") in the overall json hierarchy
+        #' @return NULL
+        append_from_json_string = function(json_string, json_forest_label) {
+            forest_container_append_from_json_string_cpp(self$forest_container_ptr, json_string, json_forest_label)
+        }, 
+
         #' @description
         #' Predict every tree ensemble on every sample in `forest_dataset`
         #' @param forest_dataset `ForestDataset` R class
@@ -106,6 +134,26 @@ ForestSamples <- R6::R6Class(
             }
         }, 
 
+        #' @description
+        #' Set a constant predicted value for every tree in the ensemble. 
+        #' Stops program if any tree is more than a root node. 
+        #' @param dataset `ForestDataset` Dataset class (covariates, basis, etc...)
+        #' @param outcome `Outcome` Outcome class (residual / partial residual)
+        #' @param forest_model `ForestModel` object storing tracking structures used in training / sampling
+        #' @param leaf_model_int Integer value encoding the leaf model type (0 = constant gaussian, 1 = univariate gaussian, 2 = multivariate gaussian, 3 = log linear variance).
+        #' @param leaf_value Constant leaf value(s) to be fixed for each tree in the ensemble indexed by `forest_num`. Can be either a single number or a vector, depending on the forest's leaf dimension.
+        prepare_for_sampler = function(dataset, outcome, forest_model, leaf_model_int, leaf_value) {
+            stopifnot(!is.null(dataset$data_ptr))
+            stopifnot(!is.null(outcome$data_ptr))
+            stopifnot(!is.null(forest_model$tracker_ptr))
+            stopifnot(!is.null(self$forest_container_ptr))
+            stopifnot(num_samples_forest_container_cpp(self$forest_container_ptr) == 0)
+
+            # Initialize the model
+            initialize_forest_model_cpp(dataset$data_ptr, outcome$data_ptr, self$forest_container_ptr, 
+                                        forest_model$tracker_ptr, leaf_value, leaf_model_int)
+        }, 
+
         #' @description
         #' Adjusts residual based on the predictions of a forest 
         #' 
@@ -294,11 +342,12 @@ ForestSamples <- R6::R6Class(
 #' @param num_trees Number of trees
 #' @param output_dimension Dimensionality of the outcome model
 #' @param is_leaf_constant Whether leaf is constant
+#' @param is_exponentiated Whether forest predictions should be exponentiated before being returned
 #'
 #' @return `ForestSamples` object
 #' @export
-createForestContainer <- function(num_trees, output_dimension=1, is_leaf_constant=F) {
+createForestContainer <- function(num_trees, output_dimension=1, is_leaf_constant=F, is_exponentiated=F) {
     return(invisible((
-        ForestSamples$new(num_trees, output_dimension, is_leaf_constant)
+        ForestSamples$new(num_trees, output_dimension, is_leaf_constant, is_exponentiated)
     )))
 }

R/kernel.R

@@ -132,17 +132,30 @@ createForestKernel <- function() {
 #' corresponds to the observations for which outcomes are unobserved and must be estimated 
 #' based on the kernels k(X_test,X_test), k(X_test,X_train), and k(X_train,X_train). If not provided, 
 #' this function will only compute k(X_train, X_train).
-#' @param forest_num (Option) Index of the forest sample to use for kernel computation. If not provided, 
+#' @param forest_num (Optional) Index of the forest sample to use for kernel computation. If not provided, 
 #' this function will use the last forest.
+#' @param forest_type (Optional) Whether to compute the kernel from the mean or variance forest. Default: "mean". Specify "variance" for the variance forest. 
+#' All other inputs are invalid. Must have sampled the relevant forest or an error will occur.
 #' @return List of kernel matrices. If `X_test = NULL`, the list contains 
 #' one `n_train` x `n_train` matrix, where `n_train = nrow(X_train)`. 
 #' This matrix is the kernel defined by `W_train %*% t(W_train)` where `W_train` 
 #' is a matrix with `n_train` rows and as many columns as there are total leaves in an ensemble. 
 #' If `X_test` is not `NULL`, the list contains two more matrices defined by 
 #' `W_test %*% t(W_train)` and `W_test %*% t(W_test)`.
 #' @export
-computeForestKernels <- function(bart_model, X_train, X_test=NULL, forest_num=NULL) {
+computeForestKernels <- function(bart_model, X_train, X_test=NULL, forest_num=NULL, forest_type="mean") {
     stopifnot(class(bart_model)=="bartmodel")
+    if (forest_type=="mean") {
+        if (!bart_model$model_params$include_mean_forest) {
+            stop("Mean forest was not sampled in the bart model provided")
+        }
+    } else if (forest_type=="variance") {
+        if (!bart_model$model_params$include_variance_forest) {
+            stop("Variance forest was not sampled in the bart model provided")
+        }
+    } else {
+        stop("Must provide either 'mean' or 'variance' for the `forest_type` parameter")
+    }
 
     # Preprocess covariates
     if (!is.data.frame(X_train)) {
@@ -164,10 +177,17 @@ computeForestKernels <- function(bart_model, X_train, X_test=NULL, forest_num=NU
     num_samples <- bart_model$model_params$num_samples
     stopifnot(forest_num <= num_samples)
     sample_index <- ifelse(is.null(forest_num), num_samples-1, forest_num-1)
-    return(forest_kernel$compute_kernel(
-        covariates_train = X_train, covariates_test = X_test,
-        forest_container = bart_model$forests, forest_num = sample_index
-    ))
+    if (forest_type=="mean") {
+        return(forest_kernel$compute_kernel(
+            covariates_train = X_train, covariates_test = X_test,
+            forest_container = bart_model$mean_forests, forest_num = sample_index
+        ))
+    } else if (forest_type=="variance") {
+        return(forest_kernel$compute_kernel(
+            covariates_train = X_train, covariates_test = X_test,
+            forest_container = bart_model$variance_forests, forest_num = sample_index
+        ))
+    }
 }
 
 #' Compute and return a vector representation of a forest's leaf predictions for 
@@ -192,21 +212,41 @@ computeForestKernels <- function(bart_model, X_train, X_test=NULL, forest_num=NU
 #' corresponds to the observations for which outcomes are unobserved and must be estimated 
 #' based on the kernels k(X_test,X_test), k(X_test,X_train), and k(X_train,X_train). If not provided, 
 #' this function will only compute k(X_train, X_train).
-#' @param forest_num (Option) Index of the forest sample to use for kernel computation. If not provided, 
+#' @param forest_num (Optional) Index of the forest sample to use for kernel computation. If not provided, 
 #' this function will use the last forest.
+#' @param forest_type (Optional) Whether to compute the kernel from the mean or variance forest. Default: "mean". Specify "variance" for the variance forest. 
+#' All other inputs are invalid. Must have sampled the relevant forest or an error will occur.
 #' @return List of vectors. If `X_test = NULL`, the list contains 
 #' one vector of length `n_train * num_trees`, where `n_train = nrow(X_train)` 
 #' and `num_trees` is the number of trees in `bart_model`. If `X_test` is not `NULL`, 
 #' the list contains another vector of length `n_test * num_trees`.
 #' @export
-computeForestLeafIndices <- function(bart_model, X_train, X_test=NULL, forest_num=NULL) {
+computeForestLeafIndices <- function(bart_model, X_train, X_test=NULL, forest_num=NULL, forest_type="mean") {
     stopifnot(class(bart_model)=="bartmodel")
+    if (forest_type=="mean") {
+        if (!bart_model$model_params$include_mean_forest) {
+            stop("Mean forest was not sampled in the bart model provided")
+        }
+    } else if (forest_type=="variance") {
+        if (!bart_model$model_params$include_variance_forest) {
+            stop("Variance forest was not sampled in the bart model provided")
+        }
+    } else {
+        stop("Must provide either 'mean' or 'variance' for the `forest_type` parameter")
+    }
     forest_kernel <- createForestKernel()
     num_samples <- bart_model$model_params$num_samples
     stopifnot(forest_num <= num_samples)
     sample_index <- ifelse(is.null(forest_num), num_samples-1, forest_num-1)
-    return(forest_kernel$compute_leaf_indices(
-        covariates_train = X_train, covariates_test = X_test,
-        forest_container = bart_model$forests, forest_num = sample_index
-    ))
+    if (forest_type == "mean") {
+        return(forest_kernel$compute_leaf_indices(
+            covariates_train = X_train, covariates_test = X_test,
+            forest_container = bart_model$mean_forests, forest_num = sample_index
+        ))
+    } else if (forest_type == "variance") {
+        return(forest_kernel$compute_leaf_indices(
+            covariates_train = X_train, covariates_test = X_test,
+            forest_container = bart_model$variance_forests, forest_num = sample_index
+        ))
+    }
 }