Deploying to r-dev from @ 022c022 🚀

Author: andrewherren

NEWS.md

@@ -3,6 +3,7 @@
 ## New Features
 
 * Replaced C++ standard library distributions (`discrete_distribution`, `uniform_real_distribution`, `normal_distribution`, and `gamma_distribution`) with custom implementations for cross-platform reproducibility.
+* Substituted custom implementations for base R `mean()`, `var()`, and `sd()` in the preprocessing logic of the R `bart()` and `bcf()` functions for enhanced numeric stability across platforms.
 
 ## Bug Fixes
 

R/bart.R

@@ -854,11 +854,11 @@ bart <- function(
   # differently for binary and continuous outcomes
   if (probit_outcome_model) {
     # Compute a probit-scale offset and fix scale to 1
-    y_bar_train <- qnorm(mean(y_train))
+    y_bar_train <- qnorm(mean_cpp(as.numeric(y_train)))
     y_std_train <- 1
 
-    # Set a pseudo outcome by subtracting mean(y_train) from y_train
-    resid_train <- y_train - mean(y_train)
+    # Set a pseudo outcome by subtracting mean_cpp(y_train) from y_train
+    resid_train <- y_train - mean_cpp(as.numeric(y_train))
 
     # Set initial values of root nodes to 0.0 (in probit scale)
     init_val_mean <- 0.0
@@ -910,8 +910,8 @@ bart <- function(
   } else {
     # Only standardize if user requested
     if (standardize) {
-      y_bar_train <- mean(y_train)
-      y_std_train <- sd(y_train)
+      y_bar_train <- mean_cpp(as.numeric(y_train))
+      y_std_train <- sd_cpp(as.numeric(y_train))
     } else {
       y_bar_train <- 0
       y_std_train <- 1
@@ -921,23 +921,23 @@ bart <- function(
     resid_train <- (y_train - y_bar_train) / y_std_train
 
     # Compute initial value of root nodes in mean forest
-    init_val_mean <- mean(resid_train)
+    init_val_mean <- mean_cpp(as.numeric(resid_train))
 
     # Calibrate priors for sigma^2 and tau
     if (is.null(sigma2_init)) {
-      sigma2_init <- 1.0 * var(resid_train)
+      sigma2_init <- 1.0 * var_cpp(as.numeric(resid_train))
     }
     if (is.null(variance_forest_init)) {
-      variance_forest_init <- 1.0 * var(resid_train)
+      variance_forest_init <- 1.0 * var_cpp(as.numeric(resid_train))
     }
     if (is.null(b_leaf)) {
-      b_leaf <- var(resid_train) / (2 * num_trees_mean)
+      b_leaf <- var_cpp(as.numeric(resid_train)) / (2 * num_trees_mean)
     }
     if (has_basis) {
       if (ncol(leaf_basis_train) > 1) {
         if (is.null(sigma2_leaf_init)) {
           sigma2_leaf_init <- diag(
-            2 * var(resid_train) / (num_trees_mean),
+            2 * var_cpp(as.numeric(resid_train)) / (num_trees_mean),
             ncol(leaf_basis_train)
           )
         }
@@ -952,7 +952,7 @@ bart <- function(
       } else {
         if (is.null(sigma2_leaf_init)) {
           sigma2_leaf_init <- as.matrix(
-            2 * var(resid_train) / (num_trees_mean)
+            2 * var_cpp(as.numeric(resid_train)) / (num_trees_mean)
           )
         }
         if (!is.matrix(sigma2_leaf_init)) {
@@ -964,7 +964,7 @@ bart <- function(
     } else {
       if (is.null(sigma2_leaf_init)) {
         sigma2_leaf_init <- as.matrix(
-          2 * var(resid_train) / (num_trees_mean)
+          2 * var_cpp(as.numeric(resid_train)) / (num_trees_mean)
         )
       }
       if (!is.matrix(sigma2_leaf_init)) {

R/bcf.R

@@ -1207,11 +1207,11 @@ bcf <- function(
   # differently for binary and continuous outcomes
   if (probit_outcome_model) {
     # Compute a probit-scale offset and fix scale to 1
-    y_bar_train <- qnorm(mean(y_train))
+    y_bar_train <- qnorm(mean_cpp(as.numeric(y_train)))
     y_std_train <- 1
 
-    # Set a pseudo outcome by subtracting mean(y_train) from y_train
-    resid_train <- y_train - mean(y_train)
+    # Set a pseudo outcome by subtracting mean_cpp(y_train) from y_train
+    resid_train <- y_train - mean_cpp(as.numeric(y_train))
 
     # Set initial value for the mu forest
     init_mu <- 0.0
@@ -1274,8 +1274,8 @@ bcf <- function(
   } else {
     # Only standardize if user requested
     if (standardize) {
-      y_bar_train <- mean(y_train)
-      y_std_train <- sd(y_train)
+      y_bar_train <- mean_cpp(as.numeric(y_train))
+      y_std_train <- sd_cpp(as.numeric(y_train))
     } else {
       y_bar_train <- 0
       y_std_train <- 1
@@ -1285,23 +1285,23 @@ bcf <- function(
     resid_train <- (y_train - y_bar_train) / y_std_train
 
     # Set initial value for the mu forest
-    init_mu <- mean(resid_train)
+    init_mu <- mean_cpp(as.numeric(resid_train))
 
     # Calibrate priors for global sigma^2 and sigma2_leaf_mu / sigma2_leaf_tau
     if (is.null(sigma2_init)) {
-      sigma2_init <- 1.0 * var(resid_train)
+      sigma2_init <- 1.0 * var_cpp(as.numeric(resid_train))
     }
     if (is.null(variance_forest_init)) {
-      variance_forest_init <- 1.0 * var(resid_train)
+      variance_forest_init <- 1.0 * var_cpp(as.numeric(resid_train))
     }
     if (is.null(b_leaf_mu)) {
-      b_leaf_mu <- var(resid_train) / (num_trees_mu)
+      b_leaf_mu <- var_cpp(as.numeric(resid_train)) / (num_trees_mu)
     }
     if (is.null(b_leaf_tau)) {
-      b_leaf_tau <- var(resid_train) / (2 * num_trees_tau)
+      b_leaf_tau <- var_cpp(as.numeric(resid_train)) / (2 * num_trees_tau)
     }
     if (is.null(sigma2_leaf_mu)) {
-      sigma2_leaf_mu <- 2.0 * var(resid_train) / (num_trees_mu)
+      sigma2_leaf_mu <- 2.0 * var_cpp(as.numeric(resid_train)) / (num_trees_mu)
       current_leaf_scale_mu <- as.matrix(sigma2_leaf_mu)
     } else {
       if (!is.matrix(sigma2_leaf_mu)) {
@@ -1311,7 +1311,7 @@ bcf <- function(
       }
     }
     if (is.null(sigma2_leaf_tau)) {
-      sigma2_leaf_tau <- var(resid_train) / (num_trees_tau)
+      sigma2_leaf_tau <- var_cpp(as.numeric(resid_train)) / (num_trees_tau)
       current_leaf_scale_tau <- as.matrix(diag(
         sigma2_leaf_tau,
         ncol(Z_train)

R/cpp11.R

@@ -276,6 +276,22 @@ root_reset_rfx_tracker_cpp <- function(tracker, dataset, residual, rfx_model) {
   invisible(.Call(`_stochtree_root_reset_rfx_tracker_cpp`, tracker, dataset, residual, rfx_model))
 }
 
+sum_cpp <- function(x) {
+  .Call(`_stochtree_sum_cpp`, x)
+}
+
+mean_cpp <- function(x) {
+  .Call(`_stochtree_mean_cpp`, x)
+}
+
+var_cpp <- function(x) {
+  .Call(`_stochtree_var_cpp`, x)
+}
+
+sd_cpp <- function(x) {
+  .Call(`_stochtree_sd_cpp`, x)
+}
+
 active_forest_cpp <- function(num_trees, output_dimension, is_leaf_constant, is_exponentiated) {
   .Call(`_stochtree_active_forest_cpp`, num_trees, output_dimension, is_leaf_constant, is_exponentiated)
 }

src/Makevars.in

@@ -22,6 +22,7 @@ OBJECTS = \
     kernel.o \
     R_data.o \
     R_random_effects.o \
+    R_utils.o \
     sampler.o \
     serialization.o \
     cpp11.o \

src/Makevars.win.in

@@ -23,6 +23,7 @@ OBJECTS = \
     kernel.o \
     R_data.o \
     R_random_effects.o \
+    R_utils.o \
     sampler.o \
     serialization.o \
     cpp11.o \

src/R_utils.cpp

@@ -0,0 +1,35 @@
+#include <cpp11.hpp>
+#include <cmath>
+
+[[cpp11::register]]
+double sum_cpp(cpp11::doubles x) {
+    double output = 0.0;
+    for (int i = 0; i < x.size(); i++) {
+        output += x[i];
+    }
+    return output;
+}
+
+[[cpp11::register]]
+double mean_cpp(cpp11::doubles x) {
+    double output = 0.0;
+    for (int i = 0; i < x.size(); i++) {
+        output += x[i];
+    }
+    return output / x.size();
+}
+
+[[cpp11::register]]
+double var_cpp(cpp11::doubles x) {
+    double mean = mean_cpp(x);
+    double output = 0.0;
+    for (int i = 0; i < x.size(); i++) {
+        output += (x[i] - mean) * (x[i] - mean);
+    }
+    return output / (x.size() - 1);
+}
+
+[[cpp11::register]]
+double sd_cpp(cpp11::doubles x) {
+    return std::sqrt(var_cpp(x));
+}

src/cpp11.cpp

@@ -516,6 +516,34 @@ extern "C" SEXP _stochtree_root_reset_rfx_tracker_cpp(SEXP tracker, SEXP dataset
     return R_NilValue;
   END_CPP11
 }
+// R_utils.cpp
+double sum_cpp(cpp11::doubles x);
+extern "C" SEXP _stochtree_sum_cpp(SEXP x) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(sum_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(x)));
+  END_CPP11
+}
+// R_utils.cpp
+double mean_cpp(cpp11::doubles x);
+extern "C" SEXP _stochtree_mean_cpp(SEXP x) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(mean_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(x)));
+  END_CPP11
+}
+// R_utils.cpp
+double var_cpp(cpp11::doubles x);
+extern "C" SEXP _stochtree_var_cpp(SEXP x) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(var_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(x)));
+  END_CPP11
+}
+// R_utils.cpp
+double sd_cpp(cpp11::doubles x);
+extern "C" SEXP _stochtree_sd_cpp(SEXP x) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(sd_cpp(cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(x)));
+  END_CPP11
+}
 // forest.cpp
 cpp11::external_pointer<StochTree::TreeEnsemble> active_forest_cpp(int num_trees, int output_dimension, bool is_leaf_constant, bool is_exponentiated);
 extern "C" SEXP _stochtree_active_forest_cpp(SEXP num_trees, SEXP output_dimension, SEXP is_leaf_constant, SEXP is_exponentiated) {
@@ -1700,6 +1728,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_stochtree_leaf_values_forest_container_cpp",                    (DL_FUNC) &_stochtree_leaf_values_forest_container_cpp,                     4},
     {"_stochtree_leaves_forest_container_cpp",                         (DL_FUNC) &_stochtree_leaves_forest_container_cpp,                          3},
     {"_stochtree_left_child_node_forest_container_cpp",                (DL_FUNC) &_stochtree_left_child_node_forest_container_cpp,                 4},
+    {"_stochtree_mean_cpp",                                            (DL_FUNC) &_stochtree_mean_cpp,                                             1},
     {"_stochtree_multiply_forest_forest_container_cpp",                (DL_FUNC) &_stochtree_multiply_forest_forest_container_cpp,                 3},
     {"_stochtree_node_depth_forest_container_cpp",                     (DL_FUNC) &_stochtree_node_depth_forest_container_cpp,                      4},
     {"_stochtree_nodes_forest_container_cpp",                          (DL_FUNC) &_stochtree_nodes_forest_container_cpp,                           3},
@@ -1781,6 +1810,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_stochtree_sample_sigma2_one_iteration_cpp",                     (DL_FUNC) &_stochtree_sample_sigma2_one_iteration_cpp,                      5},
     {"_stochtree_sample_tau_one_iteration_cpp",                        (DL_FUNC) &_stochtree_sample_tau_one_iteration_cpp,                         4},
     {"_stochtree_sample_without_replacement_integer_cpp",              (DL_FUNC) &_stochtree_sample_without_replacement_integer_cpp,               3},
+    {"_stochtree_sd_cpp",                                              (DL_FUNC) &_stochtree_sd_cpp,                                               1},
     {"_stochtree_set_leaf_value_active_forest_cpp",                    (DL_FUNC) &_stochtree_set_leaf_value_active_forest_cpp,                     2},
     {"_stochtree_set_leaf_value_forest_container_cpp",                 (DL_FUNC) &_stochtree_set_leaf_value_forest_container_cpp,                  2},
     {"_stochtree_set_leaf_vector_active_forest_cpp",                   (DL_FUNC) &_stochtree_set_leaf_vector_active_forest_cpp,                    2},
@@ -1789,12 +1819,14 @@ static const R_CallMethodDef CallEntries[] = {
     {"_stochtree_split_index_forest_container_cpp",                    (DL_FUNC) &_stochtree_split_index_forest_container_cpp,                     4},
     {"_stochtree_split_theshold_forest_container_cpp",                 (DL_FUNC) &_stochtree_split_theshold_forest_container_cpp,                  4},
     {"_stochtree_subtract_from_column_vector_cpp",                     (DL_FUNC) &_stochtree_subtract_from_column_vector_cpp,                      2},
+    {"_stochtree_sum_cpp",                                             (DL_FUNC) &_stochtree_sum_cpp,                                              1},
     {"_stochtree_sum_leaves_squared_ensemble_forest_container_cpp",    (DL_FUNC) &_stochtree_sum_leaves_squared_ensemble_forest_container_cpp,     2},
     {"_stochtree_tree_prior_cpp",                                      (DL_FUNC) &_stochtree_tree_prior_cpp,                                       4},
     {"_stochtree_update_alpha_tree_prior_cpp",                         (DL_FUNC) &_stochtree_update_alpha_tree_prior_cpp,                          2},
     {"_stochtree_update_beta_tree_prior_cpp",                          (DL_FUNC) &_stochtree_update_beta_tree_prior_cpp,                           2},
     {"_stochtree_update_max_depth_tree_prior_cpp",                     (DL_FUNC) &_stochtree_update_max_depth_tree_prior_cpp,                      2},
     {"_stochtree_update_min_samples_leaf_tree_prior_cpp",              (DL_FUNC) &_stochtree_update_min_samples_leaf_tree_prior_cpp,               2},
+    {"_stochtree_var_cpp",                                             (DL_FUNC) &_stochtree_var_cpp,                                              1},
     {NULL, NULL, 0}
 };
 }