Reformatting R and Python code

Author: andrewherren

R/bart.R

@@ -1971,9 +1971,7 @@ predict.bartmodel <- function(
   if ((predict_rfx) && (is.null(rfx_basis)) && (!rfx_intercept)) {
     stop("Random effects basis (rfx_basis) must be provided for this model")
   }
-  if (
-    (object$model_params$num_rfx_basis > 0) && (!rfx_intercept)
-  ) {
+  if ((object$model_params$num_rfx_basis > 0) && (!rfx_intercept)) {
     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"
@@ -2020,7 +2018,7 @@ predict.bartmodel <- function(
       }
     }
   }
-  
+
   # Create prediction dataset
   if (!is.null(leaf_basis)) {
     prediction_dataset <- createForestDataset(covariates, leaf_basis)
@@ -2072,7 +2070,9 @@ predict.bartmodel <- function(
     } else {
       # Sanity check -- this branch should only occur if rfx_model_spec == "intercept_only"
       if (!rfx_intercept) {
-        stop("rfx_basis must be provided for random effects models with random slopes")
+        stop(
+          "rfx_basis must be provided for random effects models with random slopes"
+        )
       }
 
       # Extract the raw RFX samples and scale by train set outcome standard deviation
@@ -2093,7 +2093,7 @@ predict.bartmodel <- function(
           rfx_beta_draws[, rfx_group_ids[i], ]
       }
 
-      # Intercept-only model, so the random effect prediction is simply the 
+      # Intercept-only model, so the random effect prediction is simply the
       # value of the respective group's intercept coefficient for each observation
       rfx_predictions = rfx_predictions_raw[, 1, ]
     }

stochtree/bart.py

@@ -260,7 +260,9 @@ def sample(
             "variance_prior_shape": 1.0,
             "variance_prior_scale": 1.0,
         }
-        rfx_params_updated = _preprocess_params(rfx_params_default, random_effects_params)
+        rfx_params_updated = _preprocess_params(
+            rfx_params_default, random_effects_params
+        )
 
         ### Unpack all parameter values
         # 1. General parameters
@@ -1459,7 +1461,9 @@ def sample(
                                 forest_dataset_train
                             )
                             if self.has_rfx:
-                                rfx_pred = rfx_model.predict(rfx_dataset_train, rfx_tracker)
+                                rfx_pred = rfx_model.predict(
+                                    rfx_dataset_train, rfx_tracker
+                                )
                                 outcome_pred = outcome_pred + rfx_pred
                             mu0 = outcome_pred[y_train[:, 0] == 0]
                             mu1 = outcome_pred[y_train[:, 0] == 1]
@@ -1835,8 +1839,10 @@ def predict(
             if rfx_basis.shape[0] != covariates.shape[0]:
                 raise ValueError("X and rfx_basis must have the same number of rows")
             if rfx_basis.shape[1] != self.num_rfx_basis:
-                raise ValueError("rfx_basis must have the same number of columns as the random effects basis used to sample this model")
-        
+                raise ValueError(
+                    "rfx_basis must have the same number of columns as the random effects basis used to sample this model"
+                )
+
         # Random effects predictions
         if predict_rfx or predict_rfx_intermediate:
             if rfx_basis is not None:
@@ -1849,10 +1855,10 @@ def predict(
                     raise ValueError(
                         "rfx_basis must be provided for random effects models with random slopes"
                     )
-                
+
                 # Extract the raw RFX samples and scale by train set outcome standard deviation
                 rfx_samples_raw = self.rfx_container.extract_parameter_samples()
-                rfx_beta_draws = rfx_samples_raw['beta_samples'] * self.y_std
+                rfx_beta_draws = rfx_samples_raw["beta_samples"] * self.y_std
 
                 # Construct an array with the appropriate group random effects arranged for each observation
                 n_train = covariates.shape[0]
@@ -1861,13 +1867,15 @@ def predict(
                         "BART models fit with random intercept models should only yield 2 dimensional random effect sample matrices"
                     )
                 else:
-                    rfx_predictions_raw = np.empty(shape=(n_train, 1, rfx_beta_draws.shape[1]))
+                    rfx_predictions_raw = np.empty(
+                        shape=(n_train, 1, rfx_beta_draws.shape[1])
+                    )
                     for i in range(n_train):
                         rfx_predictions_raw[i, 0, :] = rfx_beta_draws[
                             rfx_group_ids[i], :
                         ]
                 rfx_predictions = np.squeeze(rfx_predictions_raw[:, 0, :])
-                
+
         # Combine into y hat predictions
         if probability_scale:
             if predict_y_hat and has_mean_forest and has_rfx:
@@ -2583,9 +2591,7 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
         self.probit_outcome_model = json_object_default.get_boolean(
             "probit_outcome_model"
         )
-        self.rfx_model_spec = json_object_default.get_string(
-            "rfx_model_spec"
-        )
+        self.rfx_model_spec = json_object_default.get_string("rfx_model_spec")
 
         # Unpack number of samples
         for i in range(len(json_object_list)):

stochtree/bcf.py

@@ -317,7 +317,9 @@ def sample(
             "variance_prior_shape": 1.0,
             "variance_prior_scale": 1.0,
         }
-        rfx_params_updated = _preprocess_params(rfx_params_default, random_effects_params)
+        rfx_params_updated = _preprocess_params(
+            rfx_params_default, random_effects_params
+        )
 
         ### Unpack all parameter values
         # 1. General parameters
@@ -413,8 +415,14 @@ def sample(
         # Check random effects specification
         if not isinstance(self.rfx_model_spec, str):
             raise ValueError("rfx_model_spec must be a string")
-        if self.rfx_model_spec not in ["custom", "intercept_only", "intercept_plus_treatment"]:
-            raise ValueError("type must either be 'custom', 'intercept_only', 'intercept_plus_treatment'")
+        if self.rfx_model_spec not in [
+            "custom",
+            "intercept_only",
+            "intercept_plus_treatment",
+        ]:
+            raise ValueError(
+                "type must either be 'custom', 'intercept_only', 'intercept_plus_treatment'"
+            )
 
         # Override keep_gfr if there are no MCMC samples
         if num_mcmc == 0:
@@ -2295,7 +2303,7 @@ def predict(
     ) -> Union[dict[str, np.array], np.array]:
         """Predict outcome model components (CATE function and prognostic function) as well as overall outcome for every provided observation.
         Predicted outcomes are computed as `yhat = mu_x + Z*tau_x` where mu_x is a sample of the prognostic function and tau_x is a sample of the treatment effect (CATE) function.
-        When random effects are present, they are either included in yhat additively if `rfx_model_spec == "custom"`. They are included in mu_x if `rfx_model_spec == "intercept_only"` or 
+        When random effects are present, they are either included in yhat additively if `rfx_model_spec == "custom"`. They are included in mu_x if `rfx_model_spec == "intercept_only"` or
         partially included in mu_x and partially included in tau_x `rfx_model_spec == "intercept_plus_treatment"`.
 
         Parameters
@@ -2508,9 +2516,13 @@ def predict(
                 if rfx_basis.ndim == 1:
                     rfx_basis = np.expand_dims(rfx_basis, 1)
                 if rfx_basis.shape[0] != X.shape[0]:
-                    raise ValueError("X and rfx_basis must have the same number of rows")
+                    raise ValueError(
+                        "X and rfx_basis must have the same number of rows"
+                    )
                 if rfx_basis.shape[1] != self.num_rfx_basis:
-                    raise ValueError("rfx_basis must have the same number of columns as the random effects basis used to sample this model")
+                    raise ValueError(
+                        "rfx_basis must have the same number of columns as the random effects basis used to sample this model"
+                    )
 
         # Random effects predictions
         if predict_rfx or predict_rfx_intermediate:
@@ -2522,26 +2534,28 @@ def predict(
         if predict_rfx_raw:
             # Extract the raw RFX samples and scale by train set outcome standard deviation
             rfx_samples_raw = self.rfx_container.extract_parameter_samples()
-            rfx_beta_draws = rfx_samples_raw['beta_samples'] * self.y_std
+            rfx_beta_draws = rfx_samples_raw["beta_samples"] * self.y_std
 
             # Construct an array with the appropriate group random effects arranged for each observation
             if rfx_beta_draws.ndim == 3:
-                rfx_predictions_raw = np.empty(shape=(X.shape[0], rfx_beta_draws.shape[0], rfx_beta_draws.shape[2]))
+                rfx_predictions_raw = np.empty(
+                    shape=(X.shape[0], rfx_beta_draws.shape[0], rfx_beta_draws.shape[2])
+                )
                 for i in range(X.shape[0]):
                     rfx_predictions_raw[i, :, :] = rfx_beta_draws[
                         :, rfx_group_ids[i], :
                     ]
             elif rfx_beta_draws.ndim == 2:
-                rfx_predictions_raw = np.empty(shape=(X.shape[0], 1, rfx_beta_draws.shape[1]))
+                rfx_predictions_raw = np.empty(
+                    shape=(X.shape[0], 1, rfx_beta_draws.shape[1])
+                )
                 for i in range(X.shape[0]):
-                    rfx_predictions_raw[i, 0, :] = rfx_beta_draws[
-                        rfx_group_ids[i], :
-                    ]
+                    rfx_predictions_raw[i, 0, :] = rfx_beta_draws[rfx_group_ids[i], :]
             else:
                 raise ValueError(
                     "Unexpected number of dimensions in extracted random effects samples"
                 )
-        
+
         # Add raw RFX predictions to mu and tau if warranted by the RFX model spec
         if predict_mu_forest or predict_mu_forest_intermediate:
             if rfx_intercept and predict_rfx_raw:
@@ -2553,7 +2567,6 @@ def predict(
                 tau_x = tau_x_forest + np.squeeze(rfx_predictions_raw[:, 1:, :])
             else:
                 tau_x = tau_x_forest
-        
 
         # Combine into y hat predictions
         needs_mean_term_preds = (
@@ -3308,7 +3321,9 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
         self.has_rfx = json_object_default.get_boolean("has_rfx")
         self.has_rfx_basis = json_object_default.get_boolean("has_rfx_basis")
         self.num_rfx_basis = json_object_default.get_scalar("num_rfx_basis")
-        self.multivariate_treatment = json_object_default.get_boolean("multivariate_treatment")
+        self.multivariate_treatment = json_object_default.get_boolean(
+            "multivariate_treatment"
+        )
         if self.has_rfx:
             self.rfx_container = RandomEffectsContainer()
             for i in range(len(json_object_list)):
@@ -3344,9 +3359,7 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
         self.probit_outcome_model = json_object_default.get_boolean(
             "probit_outcome_model"
         )
-        self.rfx_model_spec = json_object_default.get_string(
-            "rfx_model_spec"
-        )
+        self.rfx_model_spec = json_object_default.get_string("rfx_model_spec")
 
         # Unpack number of samples
         for i in range(len(json_object_list)):