Initial implementation of flexible random effects parameters through the BART interface

Author: andrewherren

R/bart.R

@@ -45,6 +45,12 @@
 #'   - `num_chains` How many independent MCMC chains should be sampled. If `num_mcmc = 0`, this is ignored. If `num_gfr = 0`, then each chain is run from root for `num_mcmc * keep_every + num_burnin` iterations, with `num_mcmc` samples retained. If `num_gfr > 0`, each MCMC chain will be initialized from a separate GFR ensemble, with the requirement that `num_gfr >= num_chains`. Default: `1`.
 #'   - `verbose` Whether or not to print progress during the sampling loops. Default: `FALSE`.
 #'   - `probit_outcome_model` Whether or not the outcome should be modeled as explicitly binary via a probit link. If `TRUE`, `y` must only contain the values `0` and `1`. Default: `FALSE`.
+#'   - `rfx_working_parameter_prior_mean` Prior mean for the random effects "working parameter". Default: `NULL`. Must be a vector whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a vector.
+#'   - `rfx_group_parameters_prior_mean` Prior mean for the random effects "group parameters." Default: `NULL`. Must be a vector whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a vector.
+#'   - `rfx_working_parameter_prior_cov` Prior covariance matrix for the random effects "working parameter." Default: `NULL`. Must be a square matrix whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a diagonal matrix.
+#'   - `rfx_group_parameter_prior_cov` Prior covariance matrix for the random effects "group parameters." Default: `NULL`. Must be a square matrix whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a diagonal matrix.
+#'   - `rfx_variance_prior_shape` Shape parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`.
+#'   - `rfx_variance_prior_scale` Scale parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`.
 #'
 #' @param mean_forest_params (Optional) A list of mean forest model parameters, each of which has a default value processed internally, so this argument list is optional.
 #'
@@ -118,7 +124,13 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
         variable_weights = NULL, random_seed = -1, 
         keep_burnin = FALSE, keep_gfr = FALSE, keep_every = 1, 
         num_chains = 1, verbose = FALSE, 
-        probit_outcome_model = FALSE
+        probit_outcome_model = FALSE,
+        rfx_working_parameter_prior_mean = NULL,
+        rfx_group_parameter_prior_mean = NULL,
+        rfx_working_parameter_prior_cov = NULL,
+        rfx_group_parameter_prior_cov = NULL,
+        rfx_variance_prior_shape = 1,
+        rfx_variance_prior_scale = 1
     )
     general_params_updated <- preprocessParams(
         general_params_default, general_params
@@ -168,6 +180,12 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
     num_chains <- general_params_updated$num_chains
     verbose <- general_params_updated$verbose
     probit_outcome_model <- general_params_updated$probit_outcome_model
+    rfx_working_parameter_prior_mean <- general_params_updated$rfx_working_parameter_prior_mean
+    rfx_group_parameter_prior_mean <- general_params_updated$rfx_group_parameter_prior_mean
+    rfx_working_parameter_prior_cov <- general_params_updated$rfx_working_parameter_prior_cov
+    rfx_group_parameter_prior_cov <- general_params_updated$rfx_group_parameter_prior_cov
+    rfx_variance_prior_shape <- general_params_updated$rfx_variance_prior_shape
+    rfx_variance_prior_scale <- general_params_updated$rfx_variance_prior_scale
 
     # 2. Mean forest parameters
     num_trees_mean <- mean_forest_params_updated$num_trees
@@ -672,19 +690,43 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
 
     # Random effects initialization 
     if (has_rfx) {
-        # Prior parameters
-        if (num_rfx_components == 1) {
-            alpha_init <- c(1)
-        } else if (num_rfx_components > 1) {
-            alpha_init <- c(1,rep(0,num_rfx_components-1))
+        # Prior parameters - use user-provided values or defaults
+        if (is.null(rfx_working_parameter_prior_mean)) {
+            if (num_rfx_components == 1) {
+                alpha_init <- c(1)
+            } else if (num_rfx_components > 1) {
+                alpha_init <- c(1,rep(0,num_rfx_components-1))
+            } else {
+                stop("There must be at least 1 random effect component")
+            }
         } else {
-            stop("There must be at least 1 random effect component")
+            alpha_init <- expand_dims_1d(rfx_working_parameter_prior_mean, num_rfx_components)
         }
-        xi_init <- matrix(rep(alpha_init, num_rfx_groups),num_rfx_components,num_rfx_groups)
-        sigma_alpha_init <- diag(1,num_rfx_components,num_rfx_components)
-        sigma_xi_init <- diag(1,num_rfx_components,num_rfx_components)
-        sigma_xi_shape <- 1
-        sigma_xi_scale <- 1
+        
+        if (is.null(rfx_group_parameter_prior_mean)) {
+            xi_init <- matrix(rep(alpha_init, num_rfx_groups),num_rfx_components,num_rfx_groups)
+        } else {
+            xi_init <- expand_dims_1d(rfx_group_parameter_prior_mean, num_rfx_components)
+            # If it's a vector, expand to matrix
+            if (is.vector(xi_init)) {
+                xi_init <- matrix(rep(xi_init, num_rfx_groups), num_rfx_components, num_rfx_groups)
+            }
+        }
+        
+        if (is.null(rfx_working_parameter_prior_cov)) {
+            sigma_alpha_init <- diag(1,num_rfx_components,num_rfx_components)
+        } else {
+            sigma_alpha_init <- expand_dims_2d_diag(rfx_working_parameter_prior_cov, num_rfx_components)
+        }
+        
+        if (is.null(rfx_group_parameter_prior_cov)) {
+            sigma_xi_init <- diag(1,num_rfx_components,num_rfx_components)
+        } else {
+            sigma_xi_init <- expand_dims_2d_diag(rfx_group_parameter_prior_cov, num_rfx_components)
+        }
+        
+        sigma_xi_shape <- rfx_variance_prior_shape
+        sigma_xi_scale <- rfx_variance_prior_scale
 
         # Random effects data structure and storage container
         rfx_dataset_train <- createRandomEffectsDataset(rfx_group_ids_train, rfx_basis_train)

R/utils.R

@@ -855,3 +855,86 @@ orderedCatPreprocess <- function(x_input, unique_levels, var_name = NULL) {
     }
     return(x_preprocessed)
 }
+
+#' Convert scalar input to vector of dimension `output_size`, 
+#' or check that input array is equivalent to a vector of dimension `output_size`.
+#' 
+#' @param input Input to be converted to a vector (or passed through as-is)
+#' @param output_size Intended size of the output vector
+#' @return A vector of length `output_size`
+#' @export
+expand_dims_1d <- function(input, output_size) {
+    if (length(input) == 1) {
+        output <- rep(input, output_size)
+    } else if (is.numeric(input)) {
+        if (length(input) != output_size) {
+            stop("`input` must be a 1D numpy array with `output_size` elements")
+        }
+        output <- input
+    } else {
+        stop("`input` must be either a 1D numpy array or a scalar that can be repeated `output_size` times")
+    }
+    return(output)
+}
+
+#' Ensures that input is propagated appropriately to a matrix of dimension `output_rows` x `output_cols`. 
+#' Handles the following cases:
+#'  1. `input` is a scalar: output is simply a (`output_rows`, `output_cols`) matrix with `input` repeated for each element
+#'  2. `input` is a vector of length `output_rows`: output is a (`output_rows`, `output_cols`) array with `input` broadcast across each of `output_cols` columns
+#'  3. `input` is a vector of length `output_cols`: output is a (`output_rows`, `output_cols`) array with `input` broadcast across each of `output_rows` rows
+#'  4. `input` is a matrix of dimension (`output_rows`, `output_cols`): input is passed through as-is
+#' All other cases throw an error.
+#' 
+#' @param input Input to be converted to a matrix (or passed through as-is)
+#' @param output_rows Intended number of rows in the output array
+#' @param output_cols Intended number of columns in the output array
+#' @return A matrix of dimension `output_rows` x `output_cols`
+#' @export
+expand_dims_2d <- function(input, output_rows, output_cols) {
+    if (length(input) == 1) {
+        output <- as.matrix(rep(input, output_rows * output_cols), ncol = output_cols)
+    } else if (is.numeric(input)) {
+        if (length(input) == output_cols) {
+            output <- matrix(rep(x, output_rows), nrow=output_rows, byrow = T)
+        } else if (length(input) == output_rows) {
+            output <- matrix(rep(x, output_cols), ncol=output_cols, byrow = F)
+        } else {
+            stop("If `input` is a vector, it must either contain `output_rows` or `output_cols` elements")
+        }
+    } else if (is.matrix(input)) {
+        if (nrow(input) != output_rows) {
+            stop("`input` must be a matrix with `output_rows` rows")
+        }
+        if (ncol(input) != output_cols) {
+            stop("`input` must be a matrix with `output_cols` columns")
+        }
+        output <- input
+    } else {
+        stop("`input` must be either a matrix, vector or a scalar")
+    }
+    return(output)
+}
+
+#' Convert scalar input to square matrix of dimension `output_size` x `output_size` with `input` along the diagonal, 
+#' or check that input array is equivalent to a square matrix of dimension `output_size` x `output_size`.
+#' 
+#' @param input Input to be converted to a square matrix (or passed through as-is)
+#' @param output_size Intended row and column dimension of the square output matrix
+#' @return A square matrix of dimension `output_size` x `output_size`
+#' @export
+expand_dims_2d_diag <- function(input, output_size) {
+    if (length(input) == 1) {
+        output <- as.matrix(diag(input, output_size))
+    } else if (is.matrix(input)) {
+        if (nrow(input) != ncol(input)) {
+            stop("`input` must be a square matrix")
+        }
+        if (nrow(input) != output_size) {
+            stop("`input` must be a square matrix with `output_size` rows and columns")
+        }
+        output <- input
+    } else {
+        stop("`input` must be either a square matrix or a scalar")
+    }
+    return(output)
+}

stochtree/__init__.py

@@ -31,6 +31,9 @@
     _check_matrix_square,
     _standardize_array_to_list,
     _standardize_array_to_np,
+    _expand_dims_1d, 
+    _expand_dims_2d, 
+    _expand_dims_2d_diag
 )
 
 __all__ = [

stochtree/bart.py

@@ -23,7 +23,7 @@
 )
 from .sampler import RNG, ForestSampler, GlobalVarianceModel, LeafVarianceModel
 from .serialization import JSONSerializer
-from .utils import NotSampledError
+from .utils import NotSampledError, _expand_dims_1d, _expand_dims_2d_diag
 
 
 class BARTModel:
@@ -132,6 +132,12 @@ def sample(
             * `keep_every` (`int`): How many iterations of the burned-in MCMC sampler should be run before forests and parameters are retained. Defaults to `1`. Setting `keep_every = k` for some `k > 1` will "thin" the MCMC samples by retaining every `k`-th sample, rather than simply every sample. This can reduce the autocorrelation of the MCMC samples.
             * `num_chains` (`int`): How many independent MCMC chains should be sampled. If `num_mcmc = 0`, this is ignored. If `num_gfr = 0`, then each chain is run from root for `num_mcmc * keep_every + num_burnin` iterations, with `num_mcmc` samples retained. If `num_gfr > 0`, each MCMC chain will be initialized from a separate GFR ensemble, with the requirement that `num_gfr >= num_chains`. Defaults to `1`.
             * `probit_outcome_model` (`bool`): Whether or not the outcome should be modeled as explicitly binary via a probit link. If `True`, `y` must only contain the values `0` and `1`. Default: `False`.
+            * `rfx_working_parameter_prior_mean`: Prior mean for the random effects "working parameter". Default: `None`. Must be a 1D numpy array whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a vector.
+            * `rfx_group_parameter_prior_mean`: Prior mean for the random effects "group parameters." Default: `None`. Must be a 1D numpy array whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a vector.
+            * `rfx_working_parameter_prior_cov`: Prior covariance matrix for the random effects "working parameter." Default: `None`. Must be a square numpy matrix whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a diagonal matrix.
+            * `rfx_group_parameter_prior_cov`: Prior covariance matrix for the random effects "group parameters." Default: `None`. Must be a square numpy matrix whose dimension matches the number of random effects bases, or a scalar value that will be expanded to a diagonal matrix.
+            * `rfx_variance_prior_shape`: Shape parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`.
+            * `rfx_variance_prior_scale`: Scale parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: `1`.
 
         mean_forest_params : dict, optional
             Dictionary of mean forest model parameters, each of which has a default value processed internally, so this argument is optional.
@@ -190,6 +196,12 @@ def sample(
             "keep_every": 1,
             "num_chains": 1,
             "probit_outcome_model": False,
+            "rfx_working_parameter_prior_mean": None,
+            "rfx_group_parameter_prior_mean": None,
+            "rfx_working_parameter_prior_cov": None,
+            "rfx_group_parameter_prior_cov": None,
+            "rfx_variance_prior_shape": 1.0,
+            "rfx_variance_prior_scale": 1.0,
         }
         general_params_updated = _preprocess_params(
             general_params_default, general_params
@@ -279,6 +291,14 @@ def sample(
         drop_vars_variance = variance_forest_params_updated["drop_vars"]
         num_features_subsample_variance = variance_forest_params_updated["num_features_subsample"]
 
+        # 4. Random effects parameters
+        rfx_working_parameter_prior_mean = general_params_updated["rfx_working_parameter_prior_mean"]
+        rfx_group_parameter_prior_mean = general_params_updated["rfx_group_parameter_prior_mean"]
+        rfx_working_parameter_prior_cov = general_params_updated["rfx_working_parameter_prior_cov"]
+        rfx_group_parameter_prior_cov = general_params_updated["rfx_group_parameter_prior_cov"]
+        rfx_variance_prior_shape = general_params_updated["rfx_variance_prior_shape"]
+        rfx_variance_prior_scale = general_params_updated["rfx_variance_prior_scale"]
+
         # Override keep_gfr if there are no MCMC samples
         if num_mcmc == 0:
             keep_gfr = True
@@ -954,22 +974,43 @@ def sample(
 
         # Set up random effects structures
         if self.has_rfx:
-            if num_rfx_components == 1:
-                alpha_init = np.array([1])
-            elif num_rfx_components > 1:
-                alpha_init = np.concatenate(
-                    (
-                        np.ones(1, dtype=float),
-                        np.zeros(num_rfx_components - 1, dtype=float),
+            # Use user-provided values or defaults
+            if rfx_working_parameter_prior_mean is None:
+                if num_rfx_components == 1:
+                    alpha_init = np.array([1])
+                elif num_rfx_components > 1:
+                    alpha_init = np.concatenate(
+                        (
+                            np.ones(1, dtype=float),
+                            np.zeros(num_rfx_components - 1, dtype=float),
+                        )
                     )
-                )
+                else:
+                    raise ValueError("There must be at least 1 random effect component")
+            else:
+                alpha_init = _expand_dims_1d(rfx_working_parameter_prior_mean, num_rfx_components)
+            
+            if rfx_group_parameter_prior_mean is None:
+                xi_init = np.tile(np.expand_dims(alpha_init, 1), (1, num_rfx_groups))
+            else:
+                xi_init = _expand_dims_1d(rfx_group_parameter_prior_mean, num_rfx_components)
+                # If it's a vector, expand to matrix
+                if xi_init.ndim == 1:
+                    xi_init = np.tile(np.expand_dims(xi_init, 1), (1, num_rfx_groups))
+            
+            if rfx_working_parameter_prior_cov is None:
+                sigma_alpha_init = np.identity(num_rfx_components)
+            else:
+                sigma_alpha_init = _expand_dims_2d_diag(rfx_working_parameter_prior_cov, num_rfx_components)
+            
+            if rfx_group_parameter_prior_cov is None:
+                sigma_xi_init = np.identity(num_rfx_components)
             else:
-                raise ValueError("There must be at least 1 random effect component")
-            xi_init = np.tile(np.expand_dims(alpha_init, 1), (1, num_rfx_groups))
-            sigma_alpha_init = np.identity(num_rfx_components)
-            sigma_xi_init = np.identity(num_rfx_components)
-            sigma_xi_shape = 1.0
-            sigma_xi_scale = 1.0
+                sigma_xi_init = _expand_dims_2d_diag(rfx_group_parameter_prior_cov, num_rfx_components)
+            
+            sigma_xi_shape = rfx_variance_prior_shape
+            sigma_xi_scale = rfx_variance_prior_scale
+            
             rfx_dataset_train = RandomEffectsDataset()
             rfx_dataset_train.add_group_labels(rfx_group_ids_train)
             rfx_dataset_train.add_basis(rfx_basis_train)