first update variable split

DESCRIPTION

@@ -14,7 +14,8 @@ License: MIT + file LICENSE
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
 RoxygenNote: 7.3.1
-LinkingTo: 
+LinkingTo:
+ stochvol,
  cpp11
 Suggests: 
  knitr,
@@ -25,6 +26,7 @@ Suggests:
  mvtnorm,
  ggplot2,
  latex2exp,
+ invgamma,
  testthat (>= 3.0.0)
 VignetteBuilder: knitr
 SystemRequirements: C++17

NAMESPACE

@@ -34,7 +34,10 @@ export(oneHotEncode)
 export(oneHotInitializeAndEncode)
 export(orderedCatInitializeAndPreprocess)
 export(orderedCatPreprocess)
+export(preprocessPredictionDataFrame)
+export(preprocessTrainDataFrame)
 export(sample_sigma2_one_iteration)
 export(sample_tau_one_iteration)
 export(saveBCFModelToJsonFile)
+export(varbart)
 useDynLib(stochtree, .registration = TRUE)

R/bart.R

@@ -40,7 +40,8 @@
 #' @param random_seed Integer parameterizing the C++ random number generator. If not specified, the C++ random number generator is seeded according to `std::random_device`.
 #' @param keep_burnin Whether or not "burnin" samples should be included in cached predictions. Default FALSE. Ignored if num_mcmc = 0.
 #' @param keep_gfr Whether or not "grow-from-root" samples should be included in cached predictions. Default TRUE. Ignored if num_mcmc = 0.
-#'
+#' @param Sparse Whether you want to turn on the dirichilet prior.
+#' @param Theta_Update Whether or not update the theta of the dirichilet prior. 
 #' @return List of sampling outputs and a wrapper around the sampled forests (which can be used for in-memory prediction on new data, or serialized to JSON on disk).
 #' @export
 #'
@@ -76,7 +77,9 @@ bart <- function(X_train, y_train, W_train = NULL, group_ids_train = NULL,
  nu = 3, lambda = NULL, a_leaf = 3, b_leaf = NULL, 
  q = 0.9, sigma2_init = NULL, num_trees = 200, num_gfr = 5, 
  num_burnin = 0, num_mcmc = 100, sample_sigma = T, 
- sample_tau = T, random_seed = -1, keep_burnin = F, keep_gfr = F){
+ sample_tau = T, random_seed = -1, keep_burnin = F, keep_gfr = F,
+ Sparse = F,
+ Theta_Update = F){
  # Preprocess covariates
  if (!is.data.frame(X_train)) {
  stop("X_train must be a dataframe")
@@ -291,13 +294,17 @@ bart <- function(X_train, y_train, W_train = NULL, group_ids_train = NULL,
  # Variable selection weights
  variable_weights <- rep(1/ncol(X_train), ncol(X_train))
 
+ #Variable Selection Splits
+ variable_count_splits <- as.integer(rep(0, ncol(X_train)))
+ var_count_matrix = matrix(NA, nrow = num_samples, ncol = ncol(X_train))
+
  # Run GFR (warm start) if specified
  if (num_gfr > 0){
  gfr_indices = 1:num_gfr
  for (i in 1:num_gfr) {
  forest_model$sample_one_iteration(
  forest_dataset_train, outcome_train, forest_samples, rng, feature_types, 
- leaf_model, current_leaf_scale, variable_weights, 
+ leaf_model, current_leaf_scale, variable_weights, variable_count_splits, 
  current_sigma2, cutpoint_grid_size, gfr = T, pre_initialized = F
  )
  if (sample_sigma) {
@@ -314,6 +321,9 @@ bart <- function(X_train, y_train, W_train = NULL, group_ids_train = NULL,
  }
  }
 
+ #Dirichlet Initialization
+ theta = 1
+
  # Run MCMC
  if (num_burnin + num_mcmc > 0) {
  if (num_burnin > 0) {
@@ -323,11 +333,20 @@ bart <- function(X_train, y_train, W_train = NULL, group_ids_train = NULL,
  mcmc_indices = (num_gfr+num_burnin+1):(num_gfr+num_burnin+num_mcmc)
  }
  for (i in (num_gfr+1):num_samples) {
-   forest_model$sample_one_iteration(
+ variable_count_splits = forest_model$sample_one_iteration(
  forest_dataset_train, outcome_train, forest_samples, rng, feature_types, 
- leaf_model, current_leaf_scale, variable_weights, 
+ leaf_model, current_leaf_scale, variable_weights, variable_count_splits, 
  current_sigma2, cutpoint_grid_size, gfr = F, pre_initialized = F
  )
+ if(Sparse == TRUE){
+ lpv = draw_s(variable_count_splits, theta)
+ variable_weights = exp(lpv)
+ if(Theta_Update == TRUE){
+ theta = draw_theta0(theta, lpv, 0.5, 1, rho = length(lpv)) 
+ }
+
+ }
+ var_count_matrix[i,] = variable_count_splits
  if (sample_sigma) {
  global_var_samples[i] <- sample_sigma2_one_iteration(outcome_train, rng, nu, lambda)
  current_sigma2 <- global_var_samples[i]
@@ -339,6 +358,7 @@ bart <- function(X_train, y_train, W_train = NULL, group_ids_train = NULL,
  if (has_rfx) {
  rfx_model$sample_random_effect(rfx_dataset_train, outcome_train, rfx_tracker_train, rfx_samples, current_sigma2, rng)
  }
+
  }
  }
 
@@ -420,7 +440,8 @@ bart <- function(X_train, y_train, W_train = NULL, group_ids_train = NULL,
  "has_rfx_basis" = has_basis_rfx, 
  "num_rfx_basis" = num_basis_rfx, 
  "sample_sigma" = sample_sigma,
- "sample_tau" = sample_tau
+ "sample_tau" = sample_tau,
+ "variable_count_splits" =var_count_matrix
  )
  result <- list(
  "forests" = forest_samples, 
@@ -644,3 +665,87 @@ getRandomEffectSamples.bartmodel <- function(object, ...){
 
  return(result)
 }
+
+
+
+
+
+log_sum_exp = function(v){
+ n = length(v)
+ mx = max(v)
+ sm = 0
+ for(i in 1:n){
+ sm = sm + exp(v[i] - mx)
+ }
+ return(mx + log(sm))
+}
+
+log_gamma = function(shape){
+ y = log(rgamma(1, shape+ 1))
+ z = log(runif(1))/shape
+ return(y+z)
+}
+
+log_dirichilet = function(alpha){
+ k = length(alpha)
+ draw = rep(0,k)
+ for(j in 1:k){
+ draw[j] = log_gamma(alpha[j])
+ }
+ lse = log_sum_exp(draw)
+ for(j in 1:k){
+ draw[j] = draw[j] - lse
+ }
+ return(draw)
+}
+
+
+draw_s = function(nv,theta = 1){
+ n = length(nv)
+ theta_ = rep(0, n)
+ for(i in 1:n){
+ theta_[i] = theta/n + nv[i]
+ }
+ lpv = log_dirichilet(theta_)
+ return(lpv)
+}
+
+
+
+discrete = function(wts) {
+ p <- length(wts)
+ x <- 0
+ vOut <- rep(0, p)
+ vOut <- rmultinom(1, size = 1, prob = wts)
+ if (vOut[1] == 0) {
+ for (j in 2:p) {
+ x <- x + j * vOut[j]
+ }
+ }
+ return(x)
+}
+
+draw_theta0 = function(theta, lpv, a , b, rho) {
+ p = length(lpv)
+ sumlpv = sum(lpv)
+ lambda_g <- seq(1 / 1001, 1000 / 1001, length.out = 1000)
+ theta_g <- lambda_g * rho / (1 - lambda_g)
+ lwt_g = rep(0, 1000)
+
+ for (k in 1:1000) {
+ theta_log_lik = lgamma(theta_g[k]) - p * lgamma(theta_g[k] / p) + (theta_g[k] / p) * sumlpv
+ beta_log_prior = (a - 1) * log(lambda_g[k]) + (b - 1) * log(1 - lambda_g[k])
+ lwt_g[k] = theta_log_lik + beta_log_prior
+ }
+
+ lse <- log_sum_exp(lwt_g)
+ lwt_g <- exp(lwt_g - lse)
+ weights <- lwt_g / sum(lwt_g)
+ theta <- theta_g[discrete(weights)]
+
+ return(theta)
+}
+
+
+
+

R/bcf.R

@@ -290,6 +290,7 @@ bcf <- function(X_train, Z_train, y_train, pi_train = NULL, group_ids_train = NU
  }
 
  # Set variable weights for the prognostic and treatment effect forests
+ variable_selection_splits = NULL
  variable_weights_mu = rep(1/ncol(X_train_mu), ncol(X_train_mu))
  variable_weights_tau = rep(1/ncol(X_train_tau), ncol(X_train_tau))
 
@@ -404,7 +405,7 @@ bcf <- function(X_train, Z_train, y_train, pi_train = NULL, group_ids_train = NU
  # Sample the prognostic forest
  forest_model_mu$sample_one_iteration(
  forest_dataset_mu_train, outcome_train, forest_samples_mu, rng, feature_types_mu, 
- 0, current_leaf_scale_mu, variable_weights_mu, 
+ 0, current_leaf_scale_mu, variable_weights_mu, variable_selection_splits, 
  current_sigma2, cutpoint_grid_size, gfr = T, pre_initialized = T
  )
 
@@ -421,7 +422,7 @@ bcf <- function(X_train, Z_train, y_train, pi_train = NULL, group_ids_train = NU
  # Sample the treatment forest
  forest_model_tau$sample_one_iteration(
  forest_dataset_tau_train, outcome_train, forest_samples_tau, rng, feature_types_tau, 
- 1, current_leaf_scale_tau, variable_weights_tau, 
+ 1, current_leaf_scale_tau, variable_weights_tau, variable_selection_splits, 
  current_sigma2, cutpoint_grid_size, gfr = T, pre_initialized = T
  )
 
@@ -488,7 +489,7 @@ bcf <- function(X_train, Z_train, y_train, pi_train = NULL, group_ids_train = NU
  # Sample the prognostic forest
  forest_model_mu$sample_one_iteration(
  forest_dataset_mu_train, outcome_train, forest_samples_mu, rng, feature_types_mu, 
- 0, current_leaf_scale_mu, variable_weights_mu, 
+ 0, current_leaf_scale_mu, variable_weights_mu, variable_selection_splits, 
  current_sigma2, cutpoint_grid_size, gfr = F, pre_initialized = T
  )
 
@@ -505,7 +506,7 @@ bcf <- function(X_train, Z_train, y_train, pi_train = NULL, group_ids_train = NU
  # Sample the treatment forest
  forest_model_tau$sample_one_iteration(
  forest_dataset_tau_train, outcome_train, forest_samples_tau, rng, feature_types_tau, 
- 1, current_leaf_scale_tau, variable_weights_tau, 
+ 1, current_leaf_scale_tau, variable_weights_tau, variable_selection_splits, 
  current_sigma2, cutpoint_grid_size, gfr = F, pre_initialized = T
  )
 

R/cpp11.R

@@ -272,8 +272,8 @@ sample_gfr_one_iteration_cpp <- function(data, residual, forest_samples, tracker
  invisible(.Call(`_stochtree_sample_gfr_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized))
 }
 
-sample_mcmc_one_iteration_cpp <- function(data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized) {
- invisible(.Call(`_stochtree_sample_mcmc_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, global_variance, leaf_model_int, pre_initialized))
+sample_mcmc_one_iteration_cpp <- function(data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, variable_count_splits, global_variance, leaf_model_int, pre_initialized) {
+ .Call(`_stochtree_sample_mcmc_one_iteration_cpp`, data, residual, forest_samples, tracker, split_prior, rng, feature_types, cutpoint_grid_size, leaf_model_scale_input, variable_weights, variable_count_splits, global_variance, leaf_model_int, pre_initialized)
 }
 
 sample_sigma2_one_iteration_cpp <- function(residual, rng, nu, lambda) {

R/model.R

@@ -67,12 +67,13 @@ ForestModel <- R6::R6Class(
  #' @param leaf_model_int Integer specifying the leaf model type (0 = constant leaf, 1 = univariate leaf regression, 2 = multivariate leaf regression)
  #' @param leaf_model_scale Scale parameter used in the leaf node model (should be a q x q matrix where q is the dimensionality of the basis and is only >1 when `leaf_model_int = 2`)
  #' @param variable_weights Vector specifying sampling probability for all p covariates in `forest_dataset`
+ #' @param variable_count_splits ####
  #' @param global_scale Global variance parameter
  #' @param cutpoint_grid_size (Optional) Number of unique cutpoints to consider (default: 500, currently only used when `GFR = TRUE`)
  #' @param gfr (Optional) Whether or not the forest should be sampled using the "grow-from-root" (GFR) algorithm
  #' @param pre_initialized (Optional) Whether or not the leaves are pre-initialized outside of the sampling loop (before any samples are drawn). In multi-forest implementations like BCF, this is true, though in the single-forest supervised learning implementation, we can let C++ do the initialization. Default: F.
  sample_one_iteration = function(forest_dataset, residual, forest_samples, rng, feature_types, 
- leaf_model_int, leaf_model_scale, variable_weights, 
+ leaf_model_int, leaf_model_scale, variable_weights, variable_count_splits, 
  global_scale, cutpoint_grid_size = 500, gfr = T, 
  pre_initialized = F) {
  if (gfr) {
@@ -87,7 +88,7 @@ ForestModel <- R6::R6Class(
  forest_dataset$data_ptr, residual$data_ptr, 
  forest_samples$forest_container_ptr, self$tracker_ptr, self$tree_prior_ptr,
  rng$rng_ptr, feature_types, cutpoint_grid_size, leaf_model_scale, 
- variable_weights, global_scale, leaf_model_int, pre_initialized
+ variable_weights, variable_count_splits, global_scale, leaf_model_int, pre_initialized
  ) 
  }
  }