Merge pull request #179 from StochasticTree/make_rfx_parameters_flexible

Allow users to set random effects prior parameters in BART and BCF

Author: andrewherren

NAMESPACE

@@ -38,6 +38,9 @@ export(createRandomEffectSamples)
 export(createRandomEffectsDataset)
 export(createRandomEffectsModel)
 export(createRandomEffectsTracker)
+export(expand_dims_1d)
+export(expand_dims_2d)
+export(expand_dims_2d_diag)
 export(getRandomEffectSamples)
 export(loadForestContainerCombinedJson)
 export(loadForestContainerCombinedJsonString)

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
@@ -673,18 +691,38 @@ 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))
+        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 {
+            alpha_init <- expand_dims_1d(rfx_working_parameter_prior_mean, num_rfx_components)
+        }
+        
+        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_2d(rfx_group_parameter_prior_mean, 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 {
-            stop("There must be at least 1 random effect component")
+            sigma_xi_init <- expand_dims_2d_diag(rfx_group_parameter_prior_cov, 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
+        
+        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/bcf.R

@@ -47,6 +47,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 prognostic_forest_params (Optional) A list of prognostic forest model parameters, each of which has a default value processed internally, so this argument list is optional.
 #'
@@ -162,7 +168,13 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
         treated_coding_init = 0.5, rfx_prior_var = 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
@@ -230,6 +242,12 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
     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. Mu forest parameters
     num_trees_mu <- prognostic_forest_params_updated$num_trees
@@ -842,24 +860,39 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
 
     # Random effects prior parameters
     if (has_rfx) {
-        # Initialize the working parameter to 1
-        if (num_rfx_components < 1) {
-            stop("There must be at least 1 random effect component")
+        # Prior parameters
+        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 {
+            alpha_init <- expand_dims_1d(rfx_working_parameter_prior_mean, num_rfx_components)
+        }
+        
+        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_2d(rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups)
         }
-        alpha_init <- rep(1,num_rfx_components)
-        # Initialize each group parameter based on a regression of outcome on basis in that grou
-        xi_init <- matrix(0,num_rfx_components,num_rfx_groups)
-        for (i in 1:num_rfx_groups) {
-            group_subset_indices <- rfx_group_ids_train == i
-            basis_group <- rfx_basis_train[group_subset_indices,]
-            resid_group <- resid_train[group_subset_indices]
-            rfx_group_model <- lm(resid_group ~ 0+basis_group)
-            xi_init[,i] <- unname(coef(rfx_group_model))
+        
+        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_alpha_init <- diag(1,num_rfx_components,num_rfx_components)
-        sigma_xi_init <- diag(rfx_prior_var)
-        sigma_xi_shape <- 1
-        sigma_xi_scale <- 1
+        
+        sigma_xi_shape <- rfx_variance_prior_shape
+        sigma_xi_scale <- rfx_variance_prior_scale
     }
 
     # Random effects data structure and storage container

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 <- matrix(rep(input, output_rows * output_cols), ncol = output_cols)
+    } else if (is.numeric(input)) {
+        if (length(input) == output_cols) {
+            output <- matrix(rep(input, output_rows), nrow=output_rows, byrow = T)
+        } else if (length(input) == output_rows) {
+            output <- matrix(rep(input, 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)
+}