R/calibration.R

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

#' Calibrate the scale parameter on an inverse gamma prior for the global error variance as in Chipman et al (2022)
#' 
#' Chipman, H., George, E., Hahn, R., McCulloch, R., Pratola, M. and Sparapani, R. (2022). Bayesian Additive Regression Trees, Computational Approaches. In Wiley StatsRef: Statistics Reference Online (eds N. Balakrishnan, T. Colton, B. Everitt, W. Piegorsch, F. Ruggeri and J.L. Teugels). https://doi.org/10.1002/9781118445112.stat08288
#'
#' @param y Outcome to be modeled using BART, BCF or another nonparametric ensemble method.
#' @param X Covariates to be used to partition trees in an ensemble or series of ensemble.
#' @param W (Optional) Basis used to define a "leaf regression" model for each decision tree. The "classic" BART model assumes a constant leaf parameter, which is equivalent to a "leaf regression" on a basis of all ones, though it is not necessary to pass a vector of ones, here or to the BART function. Default: `NULL`.
#' @param nu The shape parameter for the global error variance's IG prior. The scale parameter in the Sparapani et al (2021) parameterization is defined as `nu*lambda` where `lambda` is the output of this function. Default: `3`.
#' @param quant (Optional) Quantile of the inverse gamma prior distribution represented by a linear-regression-based overestimate of `sigma^2`. Default: `0.9`.
#' @param standardize (Optional) Whether or not outcome should be standardized (`(y-mean(y))/sd(y)`) before calibration of `lambda`. Default: `TRUE`.
#'
#' @return Value of `lambda` which determines the scale parameter of the global error variance prior (`sigma^2 ~ IG(nu,nu*lambda)`)
#' @export 
#'
#' @examples
#' n <- 100
#' p <- 5
#' X <- matrix(runif(n*p), ncol = p)
#' y <- 10*X[,1] - 20*X[,2] + rnorm(n)
#' nu <- 3
#' lambda <- calibrateInverseGammaErrorVariance(y, X, nu = nu)
#' sigma2hat <- mean(resid(lm(y~X))^2)
#' mean(var(y)/rgamma(100000, nu, rate = nu*lambda) < sigma2hat)
calibrateInverseGammaErrorVariance <- function(y, X, W = NULL, nu = 3, quant = 0.9, standardize = TRUE) {
    # Compute regression basis
    if (!is.null(W)) basis <- cbind(X, W)
    else basis <- X
    # Standardize outcome if requested
    if (standardize) y <- (y-mean(y))/sd(y)
    # Compute the "regression-based" overestimate of sigma^2
    sigma2hat <- mean(resid(lm(y~basis))^2)
    # Calibrate lambda based on the implied quantile of sigma2hat
    return((sigma2hat*qgamma(1-quant,nu))/nu)
}