Added rfx model_spec to python BCF interface and included demo script to test it

Author: andrewherren

R/bcf.R

@@ -358,6 +358,19 @@ bcf <- function(
   rfx_variance_prior_shape <- rfx_params_updated$variance_prior_shape
   rfx_variance_prior_scale <- rfx_params_updated$variance_prior_scale
 
+  # Handle random effects specification
+  if (!is.character(rfx_model_spec)) {
+    stop("rfx_model_spec must be a string or character vector")
+  }
+  if (
+    !(rfx_model_spec %in%
+      c("custom", "intercept_only", "intercept_plus_treatment"))
+  ) {
+    stop(
+      "rfx_model_spec must either be 'custom', 'intercept_only', or 'intercept_plus_treatment'"
+    )
+  }
+
   # Set a function-scoped RNG if user provided a random seed
   custom_rng <- random_seed >= 0
   if (custom_rng) {
@@ -2760,9 +2773,8 @@ predict.bcfmodel <- function(
     ))
     return(NULL)
   }
-  predict_rfx_intermediate <- ((predict_y_hat && has_rfx))
-  predict_rfx_raw <- ((predict_mu_forest && has_rfx && rfx_intercept_only) ||
-    (predict_mu_forest && has_rfx && rfx_intercept_plus_treatment) ||
+  predict_rfx_intermediate <- (predict_y_hat && has_rfx)
+  predict_rfx_raw <- ((predict_mu_forest && has_rfx && rfx_intercept) ||
     (predict_tau_forest && has_rfx && rfx_intercept_plus_treatment))
   predict_mu_forest_intermediate <- (predict_y_hat && has_mu_forest)
   predict_tau_forest_intermediate <- (predict_y_hat && has_tau_forest)
@@ -2946,12 +2958,12 @@ predict.bcfmodel <- function(
 
   # Extract "raw" rfx coefficients for each rfx basis term if needed
   if (predict_rfx_raw) {
-    # Extract the raw RFX samples and scale by train set outcome sd
+    # Extract the raw RFX samples and scale by train set outcome standard deviation
     rfx_param_list <- object$rfx_samples$extract_parameter_samples()
     rfx_beta_draws <- rfx_param_list$beta_samples *
       object$model_params$outcome_scale
 
-    # Construct a matrix with the correct random effects
+    # Construct a matrix with the appropriate group random effects arranged for each observation
     rfx_predictions_raw <- array(
       NA,
       dim = c(

demo/debug/bcf_contrast_debug.py

@@ -153,6 +153,47 @@
 contrast_diff = contrast_posterior_test_comparison - contrast_posterior_test
 np.allclose(contrast_diff, 0, atol=0.001)
 
+# Now repeat the same process but via random effects model spec
+bcf_model = BCFModel()
+bcf_model.sample(
+    X_train=X_train,
+    Z_train=Z_train,
+    y_train=y_train,
+    rfx_group_ids_train=group_ids_train,
+    num_gfr=10,
+    num_burnin=0,
+    num_mcmc=1000,
+    random_effects_params={"model_spec": "intercept_plus_treatment"},
+)
+
+# Compute CATE posterior
+tau_hat_posterior_test = bcf_model.compute_contrast(
+    X_0=X_test,
+    X_1=X_test,
+    Z_0=np.zeros((n_test, 1)),
+    Z_1=np.ones((n_test, 1)),
+    rfx_group_ids_0=group_ids_test,
+    rfx_group_ids_1=group_ids_test,
+    rfx_basis_0=np.concatenate((np.ones((n_test, 1)), np.zeros((n_test, 1))), axis=1),
+    rfx_basis_1=np.ones((n_test, 2)),
+    type="posterior",
+    scale="linear",
+)
+
+# Compute the same quantity via predict
+tau_hat_posterior_test_comparison = bcf_model.predict(
+    X=X_test,
+    Z=Z_test,
+    rfx_group_ids=group_ids_test,
+    type="posterior",
+    terms="cate",
+    scale="linear",
+)
+
+# Compare results
+contrast_diff = tau_hat_posterior_test_comparison - tau_hat_posterior_test
+np.allclose(contrast_diff, 0, atol=0.001)
+
 # Generate data for a probit BCF model with random effects
 X = rng.uniform(low=0.0, high=1.0, size=(n, p))
 mu_x = X[:, 0]

src/py_stochtree.cpp

@@ -1418,6 +1418,64 @@ class RandomEffectsContainerCpp {
   int NumGroups() {
     return rfx_container_->NumGroups();
   }
+  py::array_t<double> GetBeta() {
+    int num_samples = rfx_container_->NumSamples();
+    int num_components = rfx_container_->NumComponents();
+    int num_groups = rfx_container_->NumGroups();
+    std::vector<double> beta_raw = rfx_container_->GetBeta();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({num_components, num_groups, num_samples}));
+    auto accessor = result.mutable_unchecked<3>();
+    for (int i = 0; i < num_components; i++) {
+      for (int j = 0; j < num_groups; j++) {
+        for (int k = 0; k < num_samples; k++) {
+          accessor(i,j,k) = beta_raw[k*num_groups*num_components + j*num_components + i];
+        }
+      }
+    }
+    return result;
+  }
+  py::array_t<double> GetXi() {
+    int num_samples = rfx_container_->NumSamples();
+    int num_components = rfx_container_->NumComponents();
+    int num_groups = rfx_container_->NumGroups();
+    std::vector<double> xi_raw = rfx_container_->GetXi();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({num_components, num_groups, num_samples}));
+    auto accessor = result.mutable_unchecked<3>();
+    for (int i = 0; i < num_components; i++) {
+      for (int j = 0; j < num_groups; j++) {
+        for (int k = 0; k < num_samples; k++) {
+          accessor(i,j,k) = xi_raw[k*num_groups*num_components + j*num_components + i];
+        }
+      }
+    }
+    return result;
+  }
+  py::array_t<double> GetAlpha() {
+    int num_samples = rfx_container_->NumSamples();
+    int num_components = rfx_container_->NumComponents();
+    std::vector<double> alpha_raw = rfx_container_->GetAlpha();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({num_components, num_samples}));
+    auto accessor = result.mutable_unchecked<2>();
+    for (int i = 0; i < num_components; i++) {
+      for (int j = 0; j < num_samples; j++) {
+        accessor(i,j) = alpha_raw[j*num_components + i];
+      }
+    }
+    return result;
+  }
+  py::array_t<double> GetSigma() {
+    int num_samples = rfx_container_->NumSamples();
+    int num_components = rfx_container_->NumComponents();
+    std::vector<double> sigma_raw = rfx_container_->GetSigma();
+    auto result = py::array_t<double>(py::detail::any_container<py::ssize_t>({num_components, num_samples}));
+    auto accessor = result.mutable_unchecked<2>();
+    for (int i = 0; i < num_components; i++) {
+      for (int j = 0; j < num_samples; j++) {
+        accessor(i,j) = sigma_raw[j*num_components + i];
+      }
+    }
+    return result;
+  }
   void DeleteSample(int sample_num) {
     rfx_container_->DeleteSample(sample_num);
   }
@@ -2294,6 +2352,10 @@ PYBIND11_MODULE(stochtree_cpp, m) {
     .def("NumSamples", &RandomEffectsContainerCpp::NumSamples)
     .def("NumComponents", &RandomEffectsContainerCpp::NumComponents)
     .def("NumGroups", &RandomEffectsContainerCpp::NumGroups)
+    .def("GetBeta", &RandomEffectsContainerCpp::GetBeta)
+    .def("GetXi", &RandomEffectsContainerCpp::GetXi)
+    .def("GetAlpha", &RandomEffectsContainerCpp::GetAlpha)
+    .def("GetSigma", &RandomEffectsContainerCpp::GetSigma)
     .def("DeleteSample", &RandomEffectsContainerCpp::DeleteSample)
     .def("Predict", &RandomEffectsContainerCpp::Predict)
     .def("SaveToJsonFile", &RandomEffectsContainerCpp::SaveToJsonFile)