Package index • stochtree

Supervised learning

High-level functionality for training supervised Bayesian tree ensembles (BART, XBART)

bart(): Run the BART algorithm for supervised learning.

predict(<bartmodel>): Predict from a sampled BART model on new data

compute_bart_posterior_interval(): Compute posterior credible intervals for specified terms from a fitted BART model. It supports intervals for mean functions, variance functions, random effects, and overall predictions.

compute_contrast_bart_model(): Compute a contrast using a BART model by making two sets of outcome predictions and taking their difference. This function provides the flexibility to compute any contrast of interest by specifying covariates, leaf basis, and random effects bases / IDs for both sides of a two term contrast. For simplicity, we refer to the subtrahend of the contrast as the "control" or Y0 term and the minuend of the contrast as the Y1 term, though the requested contrast need not match the "control vs treatment" terminology of a classic two-treatment causal inference problem. We mirror the function calls and terminology of the predict.bartmodel function, labeling each prediction data term with a 1 to denote its contribution to the treatment prediction of a contrast and 0 to denote inclusion in the control prediction.

sample_bart_posterior_predictive(): Sample from the posterior predictive distribution for outcomes modeled by BART

Causal inference

High-level functionality for estimating causal effects using Bayesian tree ensembles (BCF, XBCF)

bcf(): Run the Bayesian Causal Forest (BCF) algorithm for regularized causal effect estimation.

predict(<bcfmodel>): Predict from a sampled BCF model on new data

compute_bcf_posterior_interval(): Compute posterior credible intervals for BCF model terms

compute_contrast_bcf_model(): Compute a contrast using a BCF model by making two sets of outcome predictions and taking their difference. For simple BCF models with binary treatment, this will yield the same prediction as requesting terms = "cate" in the predict.bcfmodel function. For more general models, such as models with continuous / multivariate treatments or an additive random effects term with a coefficient on the treatment, this function provides the flexibility to compute a any contrast of interest by specifying covariates, treatment, and random effects bases and IDs for both sides of a two term contrast. For simplicity, we refer to the subtrahend of the contrast as the "control" or Y0 term and the minuend of the contrast as the Y1 term, though the requested contrast need not match the "control vs treatment" terminology of a classic two-arm experiment. We mirror the function calls and terminology of the predict.bcfmodel function, labeling each prediction data term with a 1 to denote its contribution to the treatment prediction of a contrast and 0 to denote inclusion in the control prediction.

sample_bcf_posterior_predictive(): Sample from the posterior predictive distribution for outcomes modeled by BCF

Low-level functionality

Serialization

Classes and functions for converting sampling artifacts to JSON and saving to disk

CppJson: Class that stores draws from an random ensemble of decision trees

createCppJson(): Create a new (empty) C++ Json object

createCppJsonFile(): Create a C++ Json object from a Json file

createCppJsonString(): Create a C++ Json object from a Json string

loadForestContainerJson(): Load a container of forest samples from json

loadForestContainerCombinedJson(): Combine multiple JSON model objects containing forests (with the same hierarchy / schema) into a single forest_container

loadForestContainerCombinedJsonString(): Combine multiple JSON strings representing model objects containing forests (with the same hierarchy / schema) into a single forest_container

loadVectorJson(): Load a vector from json

loadScalarJson(): Load a scalar from json

loadRandomEffectSamplesJson(): Load a container of random effect samples from json

loadRandomEffectSamplesCombinedJson(): Combine multiple JSON model objects containing random effects (with the same hierarchy / schema) into a single container

loadRandomEffectSamplesCombinedJsonString(): Combine multiple JSON strings representing model objects containing random effects (with the same hierarchy / schema) into a single container

saveBARTModelToJson(): Convert the persistent aspects of a BART model to (in-memory) JSON

saveBARTModelToJsonFile(): Convert the persistent aspects of a BART model to (in-memory) JSON and save to a file

saveBARTModelToJsonString(): Convert the persistent aspects of a BART model to (in-memory) JSON string

createBARTModelFromJson(): Convert an (in-memory) JSON representation of a BART model to a BART model object which can be used for prediction, etc...

createBARTModelFromJsonFile(): Convert a JSON file containing sample information on a trained BART model to a BART model object which can be used for prediction, etc...

createBARTModelFromJsonString(): Convert a JSON string containing sample information on a trained BART model to a BART model object which can be used for prediction, etc...

createBARTModelFromCombinedJson(): Convert a list of (in-memory) JSON representations of a BART model to a single combined BART model object which can be used for prediction, etc...

createBARTModelFromCombinedJsonString(): Convert a list of (in-memory) JSON strings that represent BART models to a single combined BART model object which can be used for prediction, etc...

saveBCFModelToJson(): Convert the persistent aspects of a BCF model to (in-memory) JSON

saveBCFModelToJsonFile(): Convert the persistent aspects of a BCF model to (in-memory) JSON and save to a file

saveBCFModelToJsonString(): Convert the persistent aspects of a BCF model to (in-memory) JSON string

createBCFModelFromJsonFile(): Convert a JSON file containing sample information on a trained BCF model to a BCF model object which can be used for prediction, etc...

createBCFModelFromJsonString(): Convert a JSON string containing sample information on a trained BCF model to a BCF model object which can be used for prediction, etc...

createBCFModelFromJson(): Convert an (in-memory) JSON representation of a BCF model to a BCF model object which can be used for prediction, etc...

createBCFModelFromCombinedJson(): Convert a list of (in-memory) JSON strings that represent BCF models to a single combined BCF model object which can be used for prediction, etc...

createBCFModelFromCombinedJsonString(): Convert a list of (in-memory) JSON strings that represent BCF models to a single combined BCF model object which can be used for prediction, etc...

Data

Classes and functions for preparing data for sampling algorithms

ForestDataset: Dataset used to sample a forest

createForestDataset(): Create a forest dataset object

Outcome: Outcome / partial residual used to sample an additive model.

createOutcome(): Create an outcome object

RandomEffectsDataset: Dataset used to sample a random effects model

createRandomEffectsDataset(): Create a random effects dataset object

preprocessTrainData(): Preprocess covariates. DataFrames will be preprocessed based on their column types. Matrices will be passed through assuming all columns are numeric.

preprocessPredictionData(): Preprocess covariates. DataFrames will be preprocessed based on their column types. Matrices will be passed through assuming all columns are numeric.

convertPreprocessorToJson(): Convert the persistent aspects of a covariate preprocessor to (in-memory) C++ JSON object

savePreprocessorToJsonString(): Convert the persistent aspects of a covariate preprocessor to (in-memory) JSON string

createPreprocessorFromJson(): Reload a covariate preprocessor object from a JSON string containing a serialized preprocessor

createPreprocessorFromJsonString(): Reload a covariate preprocessor object from a JSON string containing a serialized preprocessor

Forest

Classes and functions for constructing and persisting forests

Forest: Class that stores a single ensemble of decision trees (often treated as the "active forest")

createForest(): Create a forest

ForestModel: Class that defines and samples a forest model

createForestModel(): Create a forest model object

ForestSamples: Class that stores draws from an random ensemble of decision trees

createForestSamples(): Create a container of forest samples

ForestModelConfig: Object used to get / set parameters and other model configuration options for a forest model in the "low-level" stochtree interface

createForestModelConfig(): Create a forest model config object

GlobalModelConfig: Object used to get / set global parameters and other global model configuration options in the "low-level" stochtree interface

createGlobalModelConfig(): Create a global model config object

CppRNG: Class that wraps a C++ random number generator (for reproducibility)

createCppRNG(): Create an R class that wraps a C++ random number generator

calibrateInverseGammaErrorVariance(): Calibrate the scale parameter on an inverse gamma prior for the global error variance as in Chipman et al (2022)

computeForestMaxLeafIndex(): Compute and return the largest possible leaf index computable by computeForestLeafIndices for the forests in a designated forest sample container.

computeForestLeafIndices(): Compute vector of forest leaf indices

computeForestLeafVariances(): Compute vector of forest leaf scale parameters

resetActiveForest(): Reset an active forest, either from a specific forest in a ForestContainer or to an ensemble of single-node (i.e. root) trees

resetForestModel(): Re-initialize a forest model (tracking data structures) from a specific forest in a ForestContainer

Random Effects

Classes and functions for constructing and persisting random effects terms

RandomEffectSamples: Class that wraps the "persistent" aspects of a C++ random effects model (draws of the parameters and a map from the original label indices to the 0-indexed label numbers used to place group samples in memory (i.e. the first label is stored in column 0 of the sample matrix, the second label is store in column 1 of the sample matrix, etc...))

createRandomEffectSamples(): Create a RandomEffectSamples object

RandomEffectsModel: The core "model" class for sampling random effects.

createRandomEffectsModel(): Create a RandomEffectsModel object

RandomEffectsTracker: Class that defines a "tracker" for random effects models, most notably storing the data indices available in each group for quicker posterior computation and sampling of random effects terms.

createRandomEffectsTracker(): Create a RandomEffectsTracker object

getRandomEffectSamples(): Generic function for extracting random effect samples from a model object (BCF, BART, etc...)

getRandomEffectSamples(<bartmodel>): Extract raw sample values for each of the random effect parameter terms.

getRandomEffectSamples(<bcfmodel>): Extract raw sample values for each of the random effect parameter terms.

sampleGlobalErrorVarianceOneIteration(): Sample one iteration of the (inverse gamma) global variance model

sampleLeafVarianceOneIteration(): Sample one iteration of the leaf parameter variance model (only for univariate basis and constant leaf!)

resetRandomEffectsModel(): Reset a RandomEffectsModel object based on the parameters indexed by sample_num in a RandomEffectsSamples object

resetRandomEffectsTracker(): Reset a RandomEffectsTracker object based on the parameters indexed by sample_num in a RandomEffectsSamples object

rootResetRandomEffectsModel(): Reset a RandomEffectsModel object to its "default" state

rootResetRandomEffectsTracker(): Reset a RandomEffectsTracker object to its "default" state

Utilities

Miscellaneous “utility” classes and functions

sample_without_replacement(): Draw sample_size samples from population_vector without replacement, weighted by sampling_probabilities

expand_dims_1d(): Convert scalar input to vector of dimension output_size, or check that input array is equivalent to a vector of dimension output_size.

expand_dims_2d()

Ensures that input is propagated appropriately to a matrix of dimension output_rows x output_cols. Handles the following cases:

input is a scalar: output is simply a (output_rows, output_cols) matrix with input repeated for each element
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
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
input is a matrix of dimension (output_rows, output_cols): input is passed through as-is All other cases throw an error.

expand_dims_2d_diag(): 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.

Package info

High-level package details

stochtree stochtree-package: stochtree: Stochastic Tree Ensembles (XBART and BART) for Supervised Learning and Causal Inference