Updated predict method and added demo script to check the behavior

Author: andrewherren

R/bcf.R

@@ -2628,7 +2628,7 @@ bcf <- function(
 #' @param rfx_group_ids (Optional) Test set group labels used for an additive random effects model.
 #' We do not currently support (but plan to in the near future), test set evaluation for group labels
 #' that were not in the training set.
-#' @param rfx_basis (Optional) Test set basis for "random-slope" regression in additive random effects model.
+#' @param rfx_basis (Optional) Test set basis for "random-slope" regression in additive random effects model. If the model was sampled with a random effects `model_spec` of "intercept_only" or "intercept_plus_treatment", this is optional, but if it is provided, it will be used.
 #' @param type (Optional) Type of prediction to return. Options are "mean", which averages the predictions from every draw of a BCF model, and "posterior", which returns the entire matrix of posterior predictions. Default: "posterior".
 #' @param terms (Optional) Which model terms to include in the prediction. This can be a single term or a list of model terms. Options include "y_hat", "prognostic_function", "cate", "rfx", "variance_forest", or "all". If a model doesn't have random effects or variance forest predictions, but one of those terms is request, the request will simply be ignored. If none of the requested terms are present in a model, this function will return `NULL` along with a warning. Default: "all".
 #' @param scale (Optional) Scale of mean function predictions. Options are "linear", which returns predictions on the original scale of the mean forest / RFX terms, and "probability", which transforms predictions into a probability of observing `y == 1`. "probability" is only valid for models fit with a probit outcome model. Default: "linear".
@@ -2830,31 +2830,41 @@ predict.bcfmodel <- function(
     group_ids_factor <- factor(rfx_group_ids, levels = rfx_unique_group_ids)
     if (sum(is.na(group_ids_factor)) > 0) {
       stop(
-        "All random effect group labels provided in rfx_group_ids must be present in rfx_group_ids"
+        "All random effect group labels provided in rfx_group_ids must have been present at sampling time"
       )
     }
     rfx_group_ids <- as.integer(group_ids_factor)
     has_rfx <- TRUE
   }
 
   # Handle RFX model specification
-  if (object$model_params$rfx_model_spec == "custom") {
-    if (is.null(rfx_basis)) {
-      stop(
-        "A user-provided basis (`rfx_basis`) must be provided when the model was sampled with a random effects model spec set to 'custom'"
-      )
+  if (has_rfx) {
+    if (object$model_params$rfx_model_spec == "custom") {
+      if (is.null(rfx_basis)) {
+        stop(
+          "A user-provided basis (`rfx_basis`) must be provided when the model was sampled with a random effects model spec set to 'custom'"
+        )
+      }
+    } else if (object$model_params$rfx_model_spec == "intercept_only") {
+      # Only construct a basis if user-provided basis missing
+      if (is.null(rfx_basis)) {
+        rfx_basis <- matrix(
+          rep(1, nrow(X)),
+          nrow = nrow(X),
+          ncol = 1
+        )
+      }
+    } else if (
+      object$model_params$rfx_model_spec == "intercept_plus_treatment"
+    ) {
+      # Only construct a basis if user-provided basis missing
+      if (is.null(rfx_basis)) {
+        rfx_basis <- cbind(
+          rep(1, nrow(X)),
+          Z
+        )
+      }
     }
-  } else if (object$model_params$rfx_model_spec == "intercept_only") {
-    rfx_basis <- matrix(
-      rep(1, nrow(X)),
-      nrow = nrow(X),
-      ncol = 1
-    )
-  } else if (object$model_params$rfx_model_spec == "intercept_plus_treatment") {
-    rfx_basis <- cbind(
-      rep(1, nrow(X)),
-      Z
-    )
   }
 
   # Add propensities to covariate set if necessary
@@ -2953,14 +2963,18 @@ predict.bcfmodel <- function(
       rfx_predictions_raw[i, , ] <-
         rfx_beta_draws[, rfx_group_ids[i], ]
     }
+  }
 
-    # Add these RFX predictions to mu and tau if warranted by the RFX model spec
-    if (predict_mu_forest && rfx_intercept) {
+  # Add raw RFX predictions to mu and tau if warranted by the RFX model spec
+  if (predict_mu_forest || predict_mu_forest_intermediate) {
+    if (rfx_intercept && predict_rfx_raw) {
       mu_hat_final <- mu_hat_forest + rfx_predictions_raw[, 1, ]
     } else {
       mu_hat_final <- mu_hat_forest
     }
-    if (predict_tau_forest && rfx_intercept_plus_treatment) {
+  }
+  if (predict_tau_forest || predict_tau_forest_intermediate) {
+    if (rfx_intercept_plus_treatment && predict_rfx_raw) {
       tau_hat_final <- (tau_hat_forest +
         rfx_predictions_raw[, 2:ncol(rfx_basis), ])
     } else {

tools/debug/bcf_cate_debug.R

@@ -193,6 +193,61 @@ all(
   abs(tau_diff) < 0.001
 )
 
+# Now repeat the same process but via random effects model spec
+bcf_model <- bcf(
+  X_train = X_train,
+  Z_train = Z_train,
+  y_train = y_train,
+  propensity_train = pi_train,
+  rfx_group_ids_train = group_ids_train,
+  X_test = X_test,
+  Z_test = Z_test,
+  propensity_test = pi_test,
+  rfx_group_ids_test = group_ids_test,
+  num_gfr = 10,
+  num_burnin = 0,
+  num_mcmc = 1000,
+  random_effects_params = list(
+    model_spec = "intercept_plus_treatment"
+  )
+)
+
+# Compute CATE posterior
+tau_hat_posterior_test <- compute_contrast_bcf_model(
+  bcf_model,
+  X_0 = X_test,
+  X_1 = X_test,
+  Z_0 = rep(0, n_test),
+  Z_1 = rep(1, n_test),
+  propensity_0 = pi_test,
+  propensity_1 = pi_test,
+  rfx_group_ids_0 = group_ids_test,
+  rfx_group_ids_1 = group_ids_test,
+  rfx_basis_0 = cbind(1, rep(0, n_test)),
+  rfx_basis_1 = cbind(1, rep(1, n_test)),
+  type = "posterior",
+  scale = "linear"
+)
+
+# Compute the same quantity via predict
+tau_hat_posterior_test_comparison <- predict(
+  bcf_model,
+  X = X_test,
+  Z = Z_test,
+  propensity = pi_test,
+  rfx_group_ids = group_ids_test,
+  rfx_basis = rfx_basis_test,
+  type = "posterior",
+  terms = "cate",
+  scale = "linear"
+)
+
+# Compare results
+tau_diff <- tau_hat_posterior_test_comparison - tau_hat_posterior_test
+all(
+  abs(tau_diff) < 0.001
+)
+
 # Generate data for a probit BCF model with random effects
 X <- matrix(rnorm(n * p), ncol = p)
 mu_x <- X[, 1]