Merge pull request #357 from StochasticTree/cloglog-r-hotfix

Hotfix for cloglog prediction and initialization bugs

Author: andrewherren

.gitignore

@@ -3,7 +3,7 @@
 *.csv
 *.DS_Store
 lib/
-build/
+build*/
 .vscode/
 xcode/
 *.json

CHANGELOG.md

@@ -9,6 +9,7 @@
 
 ## Bug Fixes
 
+* Fixed several initialization, sampling and prediction bugs for cloglog [#357](https://github.com/StochasticTree/stochtree/pull/357)
 * Fixed num_threads pass-through bug in R BART [#339](https://github.com/StochasticTree/stochtree/pull/339)
 * Fixed probit offset bug in R and Python [#337](https://github.com/StochasticTree/stochtree/pull/337)
 * Standardized multi-chain BCF handling of internal propensity models [#334](https://github.com/StochasticTree/stochtree/pull/334)

NEWS.md

@@ -7,6 +7,7 @@
 
 ## Bug Fixes
 
+* Fixed several initialization, sampling and prediction bugs for cloglog [#357](https://github.com/StochasticTree/stochtree/pull/357)
 * Fixed num_threads pass-through bug in R BART [#339](https://github.com/StochasticTree/stochtree/pull/339)
 * Fixed probit offset bug in R and Python [#337](https://github.com/StochasticTree/stochtree/pull/337)
 * Standardized multi-chain BCF handling of internal propensity models [#334](https://github.com/StochasticTree/stochtree/pull/334)

R/bart.R

@@ -2867,8 +2867,12 @@ predict.bartmodel <- function(
           mean_forest_probabilities[, j, ] <- (1 -
             exp(
               -exp(
-                mean_forest_predictions +
-                  cloglog_cutpoint_samples[j, ]
+                sweep(
+                  mean_forest_predictions,
+                  2,
+                  cloglog_cutpoint_samples[j, ],
+                  "+"
+                )
               )
             ))
         } else if (j == cloglog_num_categories) {
@@ -2881,15 +2885,23 @@ predict.bartmodel <- function(
         } else {
           mean_forest_probabilities[, j, ] <- (exp(
             -exp(
-              mean_forest_predictions +
-                cloglog_cutpoint_samples[j - 1, ]
+              sweep(
+                mean_forest_predictions,
+                2,
+                cloglog_cutpoint_samples[j, ],
+                "+"
+              )
             )
           ) *
             (1 -
               exp(
                 -exp(
-                  mean_forest_predictions +
-                    cloglog_cutpoint_samples[j, ]
+                  sweep(
+                    mean_forest_predictions,
+                    2,
+                    cloglog_cutpoint_samples[j, ],
+                    "+"
+                  )
                 )
               )))
         }

src/forest.cpp

@@ -849,7 +849,7 @@ void initialize_forest_model_active_forest_cpp(cpp11::external_pointer<StochTree
         int n = data->NumObservations();
         std::vector<double> initial_preds(n, init_val);
         data->AddVarianceWeights(initial_preds.data(), n);
-    } else if (model_type == StochTree::ModelType::kLogLinearVariance) {
+    } else if (model_type == StochTree::ModelType::kCloglogOrdinal) {
         leaf_init_val = init_val / static_cast<double>(num_trees);
         active_forest->SetLeafValue(leaf_init_val);
         UpdateResidualEntireForest(*tracker, *data, *residual, active_forest.get(), false, std::minus<double>());

stochtree/bart.py

@@ -1699,6 +1699,9 @@ def sample(
                             residual_train,
                             True,
                         )
+                        # Undo residual update done by reconstitute_from_forest for the cloglog model
+                        if link_is_cloglog:
+                          residual_train.update_data(resid_train[:, 0])
                         # Reset leaf scale
                         if sample_sigma2_leaf:
                                 leaf_scale_double = self.leaf_scale_samples[
@@ -1754,6 +1757,9 @@ def sample(
                             residual_train,
                             True,
                         )
+                        # Undo residual update done by reconstitute_from_forest for the cloglog model
+                        if link_is_cloglog:
+                          residual_train.update_data(resid_train[:, 0])
                         # Reset leaf scale
                         if sample_sigma2_leaf and previous_leaf_var_samples is not None:
                             leaf_scale_double = previous_leaf_var_samples[
@@ -1818,6 +1824,9 @@ def sample(
                             residual_train,
                             True,
                         )
+                        # Undo residual update done by reconstitute_from_forest for the cloglog model
+                        if link_is_cloglog:
+                          residual_train.update_data(resid_train[:, 0])
                         # Reset mean forest leaf scale
                         if sample_sigma2_leaf and previous_leaf_var_samples is not None:
                             current_leaf_scale[0, 0] = sigma2_leaf