Skip to contents

Continue Sampling a BART Model

Source: R/bart.R
continueSampling.bartmodel.Rd

Continue sampling from an already-fit BART model, adding more draws of the posterior to the bartmodel object. Model terms (mean forest, variance forest, random effects, parametric terms) are initialized from their last retained sample. Training data must be passed anew to this function as the bartmodel object does not retain it. This function can only run on a model that has been sampled in session via the bart() function. Models that have been serialized to JSON can be passed directly to the bart() function to achieve a similar goal of drawing more posterior samples from a given model, though in that workflow the JSON samples are not included in the resulting bartmodel object.

The model specification interface mirrors that of bart(), but some parameters are considered fixed and thus cannot be changed in the user-provided parameter lists (i.e., global_params, mean_forest_params, variance_forest_params, rfx_params). Any model specifications that define the relevant models terms, number of trees, outcome scale and distribution, or calibration / initialization (e.g. num_trees, sample_sigma2_global, sample_sigma2_leaf, sigma2_global_init, sigma2_leaf_init, standardize, outcome_model, num_chains) cannot be changed. Specifying any of these components in a parameter list triggers a warning and will be ignored. Prior hyperparameters and sampling knobs that only impact future draws can be modified (see the parameter lists below), though this means the continued draws target a slightly different posterior. This may be desirable in some use cases, but it's important to be thoughtful about any such parameter changes.

Usage

# S3 method for class 'bartmodel'
continueSampling(
  object,
  X_train,
  y_train,
  leaf_basis_train = NULL,
  rfx_group_ids_train = NULL,
  rfx_basis_train = NULL,
  X_test = NULL,
  leaf_basis_test = NULL,
  rfx_group_ids_test = NULL,
  rfx_basis_test = NULL,
  observation_weights_train = NULL,
  observation_weights = NULL,
  num_gfr = 0,
  num_burnin = 0,
  num_mcmc = 100,
  general_params = list(),
  mean_forest_params = list(),
  variance_forest_params = list(),
  random_effects_params = list(),
  ...
)

Arguments

object

Fitted model object to continue sampling. Must be of class bartmodel and must be have been obtained by running the bart() function in session (this function is unavailable for models deserialized from JSON).

X_train

Covariates used to split trees in the ensemble. May be provided either as a dataframe or a matrix. Matrix covariates will be assumed to be all numeric. Covariates passed as a dataframe will be preprocessed based on the variable types (e.g. categorical columns stored as unordered factors will be one-hot encoded, categorical columns stored as ordered factors will passed as integers to the core algorithm, along with the metadata that the column is ordered categorical).

y_train

Outcome to be modeled by the ensemble.

leaf_basis_train

(Optional) Bases used to define a regression model y ~ W in each leaf of each regression tree. By default, BART assumes constant leaf node parameters, implicitly regressing on a constant basis of ones (i.e. y ~ 1).

rfx_group_ids_train

(Optional) Group labels used for an additive random effects model.

rfx_basis_train

(Optional) Basis for "random-slope" regression in an additive random effects model. If rfx_group_ids_train is provided with a regression basis, an intercept-only random effects model will be estimated.

X_test

(Optional) Test set of covariates used to define "out of sample" evaluation data. May be provided either as a dataframe or a matrix, but the format of X_test must be consistent with that of X_train.

leaf_basis_test

(Optional) Test set of bases used to define "out of sample" evaluation data. While a test set is optional, the structure of any provided test set must match that of the training set (i.e. if both X_train and leaf_basis_train are provided, then a test set must consist of X_test and leaf_basis_test with the same number of columns).

rfx_group_ids_test

(Optional) Test set group labels used for an additive random effects model. We do not currently support (but plan to in the near future), test set evaluation for group labels that were not in the training set.

rfx_basis_test

(Optional) Test set basis for "random-slope" regression in additive random effects model.

observation_weights_train

(Optional) Numeric vector of observation weights of length nrow(X_train). Weights are applied as y_i | - ~ N(mu(X_i), sigma^2 / w_i), so larger weights increase an observation's influence on the fit. All weights must be non-negative. Default: NULL (all observations equally weighted). Compatible with Gaussian (continuous/identity) and probit outcome models; not compatible with cloglog link functions. Note: these are referred to internally in the C++ layer as "variance weights" (var_weights), since they scale the residual variance.

observation_weights

Deprecated alias for observation_weights_train; will be removed in a future release.

num_gfr

Number of "warm-start" iterations run using the grow-from-root algorithm (He and Hahn, 2021). Default: 5.

num_burnin

Number of "burn-in" iterations of the MCMC sampler. Default: 0.

num_mcmc

Number of "retained" iterations of the MCMC sampler. Default: 100.

general_params

(Optional) A list of general (non-forest-specific) model parameters. Only the following keys can be modified by continueSampling; any other parameters specified prompt a warning and are ignored. Unspecified keys default to their value set in bart().

  • cutpoint_grid_size Maximum size of the "grid" of potential cutpoints to consider in the GFR algorithm. Default: 100.

  • sigma2_global_shape Shape parameter in the IG(sigma2_global_shape, sigma2_global_scale) global error variance model. Default: 0.

  • sigma2_global_scale Scale parameter in the IG(sigma2_global_shape, sigma2_global_scale) global error variance model. Default: 0.

  • variable_weights Numeric weights reflecting the relative probability of splitting on each variable. Does not need to sum to 1 but cannot be negative. Defaults to rep(1/ncol(X_train), ncol(X_train)) if not set here. Note that if the propensity score is included as a covariate in either forest, its weight will default to 1/ncol(X_train).

  • random_seed Integer parameterizing the C++ random number generator. If not specified, the C++ random number generator is seeded according to std::random_device.

  • keep_burnin Whether or not "burnin" samples should be included in the stored samples of forests and other parameters. Default FALSE. Ignored if num_mcmc = 0.

  • keep_gfr Whether or not "grow-from-root" samples should be included in the stored samples of forests and other parameters. Default TRUE. Ignored if num_mcmc = 0.

  • keep_every How many iterations of the burned-in MCMC sampler should be run before forests and parameters are retained. Default 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 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. Note that if num_chains > 1, the returned model object will contain samples from all chains, stored consecutively. That is, if there are 4 chains with 100 samples each, the first 100 samples will be from chain 1, the next 100 samples will be from chain 2, etc... For more detail on working with multi-chain BART models, see the multi chain vignette.

  • verbose Whether or not to print progress during the sampling loops. Default: FALSE.

  • num_threads Number of threads to use in the GFR and MCMC algorithms, as well as prediction. Defaults to 1 (single-threaded). Set to -1 to use the maximum number of available threads, or a positive integer for a specific count. OpenMP must be available for values other than 1.

mean_forest_params

(Optional) A list of mean forest model parameters. Only the following keys can be modified by continueSampling; any other parameters specified prompt a warning and are ignored. Unspecified keys default to their value set in bart().

  • alpha Prior probability of splitting for a tree of depth 0 in the mean model. Tree split prior combines alpha and beta via alpha*(1+node_depth)^-beta. Default: 0.95.

  • beta Exponent that decreases split probabilities for nodes of depth > 0 in the mean model. Tree split prior combines alpha and beta via alpha*(1+node_depth)^-beta. Default: 2.

  • min_samples_leaf Minimum allowable size of a leaf, in terms of training samples, in the mean model. Default: 5.

  • max_depth Maximum depth of any tree in the ensemble in the mean model. Default: 10. Can be overridden with -1 which does not enforce any depth limits on trees.

  • sigma2_leaf_shape Shape parameter in the IG(sigma2_leaf_shape, sigma2_leaf_scale) leaf node parameter variance model. Default: 3.

  • sigma2_leaf_scale Scale parameter in the IG(sigma2_leaf_shape, sigma2_leaf_scale) leaf node parameter variance model. Calibrated internally as 0.5/num_trees if not set here.

  • keep_vars Vector of variable names or column indices denoting variables that should be included in the forest. Default: NULL.

  • drop_vars Vector of variable names or column indices denoting variables that should be excluded from the forest. Default: NULL. If both drop_vars and keep_vars are set, drop_vars will be ignored.

  • num_features_subsample How many features to subsample when growing each tree for the GFR algorithm. Defaults to the number of features in the training dataset.

variance_forest_params

(Optional) A list of variance forest model parameters. Only the following keys can be modified by continueSampling; any other parameters specified prompt a warning and are ignored. Unspecified keys default to their value set in bart().

  • alpha Prior probability of splitting for a tree of depth 0 in the variance model. Tree split prior combines alpha and beta via alpha*(1+node_depth)^-beta. Default: 0.95.

  • beta Exponent that decreases split probabilities for nodes of depth > 0 in the variance model. Tree split prior combines alpha and beta via alpha*(1+node_depth)^-beta. Default: 2.

  • min_samples_leaf Minimum allowable size of a leaf, in terms of training samples, in the variance model. Default: 5.

  • max_depth Maximum depth of any tree in the ensemble in the variance model. Default: 10. Can be overridden with -1 which does not enforce any depth limits on trees.

  • var_forest_prior_shape Shape parameter in the IG(var_forest_prior_shape, var_forest_prior_scale) conditional error variance model (which is only sampled if num_trees > 0). Calibrated internally as num_trees / leaf_prior_calibration_param^2 + 0.5 if not set.

  • var_forest_prior_scale Scale parameter in the IG(var_forest_prior_shape, var_forest_prior_scale) conditional error variance model (which is only sampled if num_trees > 0). Calibrated internally as num_trees / leaf_prior_calibration_param^2 if not set.

  • keep_vars Vector of variable names or column indices denoting variables that should be included in the forest. Default: NULL.

  • drop_vars Vector of variable names or column indices denoting variables that should be excluded from the forest. Default: NULL. If both drop_vars and keep_vars are set, drop_vars will be ignored.

  • num_features_subsample How many features to subsample when growing each tree for the GFR algorithm. Defaults to the number of features in the training dataset.

random_effects_params

(Optional) A list of random effects model parameters. Only the following keys can be modified by continueSampling; any other parameters specified prompt a warning and are ignored. Unspecified keys default to their value set in bart().

  • variance_prior_shape Shape parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: 1.

  • variance_prior_scale Scale parameter for the inverse gamma prior on the variance of the random effects "group parameter." Default: 1.

...

Other parameters (ignored).

Value

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).

Examples

n <- 100
p <- 5
X <- matrix(runif(n*p), ncol = p)
f_XW <- (
    ((0 <= X[,1]) & (0.25 > X[,1])) * (-7.5) +
    ((0.25 <= X[,1]) & (0.5 > X[,1])) * (-2.5) +
    ((0.5 <= X[,1]) & (0.75 > X[,1])) * (2.5) +
    ((0.75 <= X[,1]) & (1 > X[,1])) * (7.5)
)
noise_sd <- 1
y <- f_XW + rnorm(n, 0, noise_sd)
test_set_pct <- 0.2
n_test <- round(test_set_pct*n)
n_train <- n - n_test
test_inds <- sort(sample(1:n, n_test, replace = FALSE))
train_inds <- (1:n)[!((1:n) %in% test_inds)]
X_test <- X[test_inds,]
X_train <- X[train_inds,]
y_test <- y[test_inds]
y_train <- y[train_inds]
bart_model <- bart(X_train = X_train, y_train = y_train,
                   num_gfr = 10, num_burnin = 0, num_mcmc = 10)
continueSampling(bart_model, X_train = X_train, y_train = y_train,
                 num_gfr = 10, num_burnin = 0, num_mcmc = 10)
#> stochtree::bart() run with mean forest, global error variance model, and mean forest leaf scale model
#> Continuous outcome was modeled as Gaussian with a constant leaf prior for the mean forest
#> Outcome was standardized
#> The sampler was run for 20 GFR iterations, with 1 chain of 0 burn-in iterations and 20 MCMC iterations, retaining every iteration (i.e. no thinning)