add standard error for interval predictions (#978)

Author: hfrick

NEWS.md

@@ -4,6 +4,9 @@
 
 * `augment()` now works for censored regression models. 
 
+* For BART models with the `dbarts` engine, `predict()` can now also return the standard error for confidence and prediction intervals (#976).
+
+
 # parsnip 1.1.0
 
 This release of parsnip contains a number of new features and bug fixes, accompanied by several optimizations that substantially decrease the time to `fit()` and `predict()` with the package.

R/bart.R

@@ -232,6 +232,9 @@ dbart_predict_calc <- function(obj, new_data, type, level = 0.95, std_err = FALS
           )
         )
     }
+    if (std_err) {
+      res$.std_error <- apply(post_dist, 2, stats::sd, na.rm = TRUE)
+    }
   }
   res
 }

R/bart_data.R

@@ -140,7 +140,8 @@ set_pred(
         obj = expr(object),
         new_data = expr(new_data),
         type = "conf_int",
-        level = expr(level)
+        level = expr(level),
+        std_err = expr(std_error)
       )
   )
 )
@@ -158,7 +159,8 @@ set_pred(
         obj = expr(object),
         new_data = expr(new_data),
         type = "pred_int",
-        level = expr(level)
+        level = expr(level),
+        std_err = expr(std_error)
       )
   )
 )
@@ -215,7 +217,8 @@ set_pred(
         obj = expr(object),
         new_data = expr(new_data),
         type = "conf_int",
-        level = expr(level)
+        level = expr(level),
+        std_err = expr(std_error)
       )
   )
 )
@@ -233,7 +236,8 @@ set_pred(
         obj = expr(object),
         new_data = expr(new_data),
         type = "pred_int",
-        level = expr(level)
+        level = expr(level),
+        std_err = expr(std_error)
       )
   )
 )
@@ -248,7 +252,9 @@ set_pred(
     post = NULL,
     func = c(pkg = "parsnip", fun = "dbart_predict_calc"),
     args =
-      list(obj = quote(object),
-           new_data =  quote(new_data))
+      list(
+        obj = quote(object),
+        new_data =  quote(new_data)
+      )
   )
 )