Fixed yhat bug in adaptive coding BCF for the R interface

Author: andrewherren

R/bcf.R

@@ -2474,6 +2474,8 @@ bcf <- function(
     )
     tau_hat_train <- t(t(tau_hat_train_raw) * (b_1_samples - b_0_samples)) *
       y_std_train
+    control_adj_train <- t(t(tau_hat_train_raw) * b_0_samples) * y_std_train
+    mu_hat_train <- mu_hat_train + control_adj_train
   } else {
     tau_hat_train <- forest_samples_tau$predict_raw(forest_dataset_train) *
       y_std_train
@@ -2508,6 +2510,8 @@ bcf <- function(
         t(tau_hat_test_raw) * (b_1_samples - b_0_samples)
       ) *
         y_std_train
+      control_adj_test <- t(t(tau_hat_test_raw) * b_0_samples) * y_std_train
+      mu_hat_test <- mu_hat_test + control_adj_test
     } else {
       tau_hat_test <- forest_samples_tau$predict_raw(
         forest_dataset_test
@@ -2849,10 +2853,11 @@ predict.bcfmodel <- function(
           "all"
         ))
     ) {
-      stop(paste0(
+      warning(paste0(
         "Term '",
         term,
-        "' was requested. Valid terms are 'y_hat', 'prognostic_function', 'mu', 'cate', 'tau', 'rfx', 'variance_forest', and 'all'."
+        "' was requested. Valid terms are 'y_hat', 'prognostic_function', 'mu', 'cate', 'tau', 'rfx', 'variance_forest', and 'all'.",
+        " This term will be ignored and prediction will only proceed if other requested terms are available in the model."
       ))
     }
   }
@@ -3056,6 +3061,8 @@ predict.bcfmodel <- function(
         t(tau_hat_raw) * (object$b_1_samples - object$b_0_samples)
       ) *
         y_std
+      control_adj <- t(t(tau_hat_raw) * object$b_0_samples) * y_std
+      mu_hat_forest <- mu_hat_forest + control_adj
     } else {
       tau_hat_forest <- object$forests_tau$predict_raw(forest_dataset_pred) *
         y_std

tools/simulations/bcf-pred-rmse.R

@@ -0,0 +1,74 @@
+# Load library
+library(stochtree)
+
+# Simulation parameters
+n <- 500
+p <- 5
+n_sim <- 100
+test_set_pct <- 0.2
+
+# Simulation container
+rmses_cached <- rep(NA_real_, n_sim)
+rmses_pred <- rep(NA_real_, n_sim)
+
+# Run the simulation
+for (i in 1:n_sim) {
+  # Generate data
+  X <- matrix(rnorm(n * p), ncol = p)
+  mu_x <- X[, 1]
+  tau_x <- 0.25 * X[, 2]
+  pi_x <- pnorm(0.5 * X[, 1])
+  Z <- rbinom(n, 1, pi_x)
+  E_XZ <- mu_x + Z * tau_x
+  snr <- 2
+  y <- E_XZ + rnorm(n, 0, 1) * (sd(E_XZ) / snr)
+
+  # Train-test split
+  n_test <- round(test_set_pct * n)
+  n_train <- n - n_test
+  test_inds <- sort(sample(1:n, n_test, replace = FALSE))
+  train_inds <- (1:n)[!((1:n) %in% test_inds)]
+  X_test <- X[test_inds, ]
+  X_train <- X[train_inds, ]
+  pi_test <- pi_x[test_inds]
+  pi_train <- pi_x[train_inds]
+  Z_test <- Z[test_inds]
+  Z_train <- Z[train_inds]
+  y_test <- y[test_inds]
+  y_train <- y[train_inds]
+  mu_test <- mu_x[test_inds]
+  mu_train <- mu_x[train_inds]
+  tau_test <- tau_x[test_inds]
+  tau_train <- tau_x[train_inds]
+  E_XZ_test <- E_XZ[test_inds]
+  E_XZ_train <- E_XZ[train_inds]
+
+  # Fit a simple BCF model
+  bcf_model <- bcf(
+    X_train = X_train,
+    Z_train = Z_train,
+    y_train = y_train,
+    propensity_train = pi_train,
+    X_test = X_test,
+    Z_test = Z_test,
+    propensity_test = pi_test
+  )
+
+  # Predict out of sample
+  y_hat_test <- predict(
+    bcf_model,
+    X = X_test,
+    Z = Z_test,
+    propensity = pi_test,
+    type = "mean",
+    terms = "y_hat"
+  )
+
+  # Compute RMSE using both cached predictions and those returned by predict()
+  rmses_cached[i] <- sqrt(mean((rowMeans(bcf_model$y_hat_test) - E_XZ_test)^2))
+  rmses_pred[i] <- sqrt(mean((y_hat_test - E_XZ_test)^2))
+}
+
+# Inspect results
+mean(rmses_cached)
+mean(rmses_pred)