StochTree: BART-based modeling in R and Python

TL;DR

StochTree is a versatile C++ library offering Bayesian tree ensemble models with R and Python bindings, supporting advanced features like heteroskedasticity, random effects, and custom model configurations.

Contribution

It introduces a comprehensive, interoperable library for Bayesian tree models with flexible fit handling and lower-level access for custom modeling in R and Python.

Findings

Supports a wide range of Bayesian tree models.

Enables saving, reinitializing, and passing fits between R and Python.

Facilitates custom MCMC routines for nonstandard models.

Abstract

stochtree is a C++ library for Bayesian tree ensemble models such as BART and Bayesian Causal Forests (BCF), as well as user-specified variations. Unlike previous BART packages, stochtree provides bindings to both R and Python for full interoperability. stochtree boasts a more comprehensive range of models relative to previous packages, including heteroskedastic forests, random effects, and treed linear models. Additionally, stochtree offers flexible handling of model fits: the ability to save model fits, reinitialize models from existing fits (facilitating improved model initialization heuristics), and pass fits between R and Python. On both platforms, stochtree exposes lower-level functionality, allowing users to specify models incorporating Bayesian tree ensembles without needing to modify C++ code. We illustrate the use of stochtree in three settings: i) straightfoward applications of existing models such as BART and BCF, ii) models that include more sophisticated components like heteroskedasticity and leaf-wise regression models, and iii) as a component of custom MCMC routines to fit nonstandard tree ensemble models.