Updated predict and demo script

Author: andrewherren

R/bcf.R

@@ -2808,15 +2808,16 @@ predict.bcfmodel <- function(
     )
   }
   if ((object$model_params$has_rfx_basis) && (is.null(rfx_basis))) {
-    stop("Random effects basis (rfx_basis) must be provided for this model")
+    if (object$model_params$rfx_model_spec == "custom") {
+      stop("Random effects basis (rfx_basis) must be provided for this model")
+    }
   }
-  if (
-    (object$model_params$num_rfx_basis > 0) &&
-      (ncol(rfx_basis) != object$model_params$num_rfx_basis)
-  ) {
-    stop(
-      "Random effects basis has a different dimension than the basis used to train this model"
-    )
+  if ((object$model_params$num_rfx_basis > 0) && (!is.null(rfx_basis))) {
+    if (ncol(rfx_basis) != object$model_params$num_rfx_basis) {
+      stop(
+        "Random effects basis has a different dimension than the basis used to train this model"
+      )
+    }
   }
 
   # Preprocess covariates

tools/debug/bcf_cate_debug.R

@@ -229,14 +229,13 @@ tau_hat_posterior_test <- compute_contrast_bcf_model(
   scale = "linear"
 )
 
-# Compute the same quantity via predict
+# Compute the same quantity directly 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"