R/postestimation.R

In hsstan: Hierarchical Shrinkage Stan Models for Biomarker Selection

##=============================================================================
##
## Copyright (c) 2019 Marco Colombo
##
## Parts of the code are based on https://avehtari.github.io/bayes_R2/bayes_R2.html
## Portions copyright (c) 2019 Aki Vehtari, Andrew Gelman, Ben Goodrich, Jonah Gabry
##
## Parts of the code are based on https://github.com/stan-dev/rstanarm
## Portions copyright (c) 2015, 2016, 2017 Trustees of Columbia University
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##
##=============================================================================


#' Pointwise log-likelihood
#'
#' Compute the pointwise log-likelihood.
#'
#' @param object An object of class `hsstan`.
#' @param newdata Optional data frame to use to evaluate the log-likelihood.
#'        If `NULL` (default), the model matrix is used.
#' @param ... Currently ignored.
#'
#' @return
#' A matrix of size `S` by `N`, where `S` is number of draws from the posterior
#' distribution, and `N` is the number of data points.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' log_lik(hs.biom)
#'
#' @importFrom rstantools log_lik
#' @importFrom stats rbinom rnorm
#' @method log_lik hsstan
#' @aliases log_lik
#' @export log_lik
#' @export
log_lik.hsstan <- function(object, newdata=NULL, ...) {
    if (is.null(newdata))
        newdata <- object$data
    y <- validate.outcome(newdata[[object$model.terms$outcome]])
    mu <- posterior_linpred(object, newdata, transform=TRUE)
    if (!is.logistic(object))
        sigma <- as.matrix(object$stanfit, pars="sigma")
    llkfun <- ifelse(is.logistic(object),
                     function(x) stats::dbinom(y[x], 1, mu[, x], log=TRUE),
                     function(x) stats::dnorm(y[x], mu[, x], sigma, log=TRUE))
    llk <- cbind(sapply(1:ncol(mu), llkfun))
    return(llk)
}

#' Posterior uncertainty intervals
#'
#' Compute posterior uncertainty intervals for `hsstan` objects.
#'
#' @param object An object of class `hsstan`.
#' @param pars Names of parameters for which posterior intervals should be
#'        returned, which can be specified as regular expressions. If `NULL`
#'        (default) then this refers to the set of predictors used in the model.
#' @param prob A value between 0 and 1 indicating the desired probability
#'        to be covered by the uncertainty intervals (0.95, by default).
#' @param ... Currently ignored.
#'
#' @return
#' A matrix with lower and upper interval bounds as columns and as many rows
#' as selected parameters.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' posterior_interval(hs.biom)
#'
#' @importFrom rstantools posterior_interval
#' @method posterior_interval hsstan
#' @aliases posterior_interval
#' @export posterior_interval
#' @export
posterior_interval.hsstan <- function(object, pars=NULL, prob=0.95, ...) {
    validate.samples(object)
    validate.probability(prob)
    pars <- get.pars(object, pars)
    post.matrix <- as.matrix(object$stanfit, pars=pars)
    rstantools::posterior_interval(post.matrix, prob=prob)
}

#' Posterior distribution of the linear predictor
#'
#' Extract the posterior draws of the linear predictor, possibly transformed
#' by the inverse-link function.
#'
#' @param object An object of class `hsstan`.
#' @param transform Whether the linear predictor should be transformed using
#'        the inverse-link function (`FALSE` by default).
#' @param newdata Optional data frame containing the variables to use to
#'        predict. If `NULL` (default), the model matrix is used. If
#'        specified, its continuous variables should be standardized, since
#'        the model coefficients are learnt on standardized data.
#' @param ... Currently ignored.
#'
#' @return
#' A matrix of size `S` by `N`, where `S` is the number of draws from the
#' posterior distribution of the (transformed) linear predictor, and `N` is
#' the number of data points.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' posterior_linpred(hs.biom)
#'
#' @importFrom rstantools posterior_linpred
#' @method posterior_linpred hsstan
#' @aliases posterior_linpred
#' @export posterior_linpred
#' @export
posterior_linpred.hsstan <- function(object, transform=FALSE,
                                     newdata=NULL, ...) {
    validate.samples(object)
    newdata <- validate.newdata(object, newdata)
    pars <- grep("^beta_", object$stanfit@model_pars, value=TRUE)
    post.matrix <- as.matrix(object$stanfit, pars=pars)
    linear.predictor <- multiplyABt(post.matrix, newdata)
    if (!transform)
        return(linear.predictor)
    return(object$family$linkinv(linear.predictor))
}

#' Posterior predictive distribution
#'
#' Draw from the posterior predictive distribution of the outcome.
#'
#' @param object An object of class `hsstan`.
#' @param newdata Optional data frame containing the variables to use to
#'        predict. If `NULL` (default), the model matrix is used. If
#'        specified, its continuous variables should be standardized, since
#'        the model coefficients are learnt on standardized data.
#' @param nsamples A positive integer indicating the number of posterior samples
#'        to use. If `NULL` (default) all samples are used.
#' @param seed Optional integer defining the seed for the pseudo-random number
#'        generator.
#' @param ... Currently ignored.
#'
#' @return
#' A matrix of size `S` by `N`, where `S` is the number of simulations from
#' the posterior predictive distribution, and `N` is the number of data points.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' posterior_predict(hs.biom)
#'
#' @importFrom rstantools posterior_predict
#' @importFrom stats rbinom rnorm
#' @method posterior_predict hsstan
#' @aliases posterior_predict
#' @export posterior_predict
#' @export
posterior_predict.hsstan <- function(object, newdata=NULL, nsamples=NULL,
                                     seed=NULL, ...) {
    validate.samples(object)
    if (!is.null(nsamples))
        validate.positive.scalar(nsamples, "nsamples", int=TRUE)

    ## extract a random subset of the posterior samples
    if (!is.null(seed))
        set.seed(seed)
    num.samples <- nsamples(object)
    samp <- sample(num.samples, min(num.samples, nsamples))
    nsamples <- length(samp)

    ## generate the posterior predictions
    mu <- posterior_linpred(object, newdata, transform=TRUE)[samp, , drop=FALSE]
    nobs <- ncol(mu)
    if (is.logistic(object))
        pp <- t(sapply(1:nsamples, function(z) rbinom(nobs, 1, mu[z, ])))
    else {
        sigma <- as.matrix(object$stanfit, pars="sigma")[samp, , drop=FALSE]
        pp <- t(sapply(1:nsamples, function(z) rnorm(nobs, mu[z, ], sigma[z])))
    }
    return(pp)
}

#' Posterior measures of performance
#'
#' Compute the log-likelihood and a relevant measure of performance (R-squared
#' or AUC) from the posterior samples.
#'
#' @param obj An object of class `hsstan` or `kfold`.
#' @param prob Width of the posterior interval (0.95, by default). It is
#'        ignored if `summary=FALSE`.
#' @param sub.idx Vector of indices of observations in the dataset to be used
#'        in computing the performance measures. If `NULL` (default), all
#'        observations in the dataset are used.
#' @param summary Whether a summary of the distribution of the performance
#'        measure should be returned rather than the pointwise values
#'        (`TRUE` by default).
#' @param cores Number of cores to use for parallelization (the value of
#'        `options("mc.cores")` by default).
#'
#' @return
#' The mean, standard deviation and posterior interval of the performance
#' measure (R-squared or AUC) if `summary=TRUE`, or a vector of values
#' of the performance measure with length equal to the size of the posterior
#' sample if `summary=FALSE`. Attribute `type` reports whether the performance
#' measures are cross-validated or not. If `sub.idx` is not `NULL`, attribute
#' `subset` reports the index of observations used in the computations.
#'
#' @examples
#' \donttest{
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' # continued from ?hsstan
#' posterior_performance(hs.biom, cores=1)
#' }
#'
#' @export
posterior_performance <- function(obj, prob=0.95, sub.idx=NULL, summary=TRUE,
                                  cores=getOption("mc.cores", 1)) {
    if (inherits(obj, "hsstan")) {
        obj <- list(fits=array(list(fit=obj, test.idx=1:nrow(obj$data)), c(1, 2)),
                    data=obj$data)
        colnames(obj$fits) <- c("fit", "test.idx")
    } else if (inherits(obj, c("kfold", "loo"))) {
        if (is.null(obj[["fits"]]))
            stop("No fitted models found, run 'kfold' with store.fits=TRUE.")
    } else
        stop("Not an 'hsstan' or 'kfold' object.")

    if (is.null(sub.idx)) {
        sub.idx <- 1:nrow(obj$data)
        used.subset <- FALSE
    } else {
        validate.indices(sub.idx, nrow(obj$data), "sub.idx")
        sub.idx <- sort(sub.idx)
        used.subset <- length(sub.idx) < nrow(obj$data)
    }

    validate.samples(obj$fits[[1]])
    validate.probability(prob)
    logistic <- is.logistic(obj$fits[[1]])
    num.folds <- nrow(obj$fits)

    ## loop over the folds
    y <- mu <- llk <- NULL
    for (fold in 1:num.folds) {
        hs <- obj$fits[[fold]]
        test.idx <- intersect(obj$fits[, "test.idx"][[fold]], sub.idx)
        if (length(test.idx) == 0)
            next
        newdata <- obj$data[test.idx, ]
        y <- c(y, obj$data[test.idx, hs$model.terms$outcome])
        mu <- cbind(mu, posterior_linpred(hs, newdata=newdata, transform=TRUE))
        llk <- cbind(llk, log_lik(hs, newdata=newdata))
    }

    if (used.subset && logistic && length(unique(y)) != 2)
        stop("'sub.idx' must contain both outcome classes.")

    if (logistic) {
        par.auc <- function(i)
            as.numeric(pROC::roc(y, mu[i, ], direction="<", quiet=TRUE)$auc)
        if (.Platform$OS.type != "windows") {
            out <- parallel::mclapply(X=1:nrow(mu), mc.cores=cores,
                                      mc.preschedule=TRUE, FUN=par.auc)
        } else { # windows
            cl <- parallel::makePSOCKcluster(cores)
            on.exit(parallel::stopCluster(cl))
            out <- parallel::parLapply(X=1:nrow(mu), cl=cl, fun=par.auc)
        }
    } else {
        out <- pmax(fastCor(y, mu), 0)^2
    }

    out <- cbind(perf=unlist(out), llk=rowSums(llk))
    colnames(out)[1] <- ifelse(logistic, "auc", "r2")
    if (summary)
        out <- posterior_summary(out, prob)
    attr(out, "type") <- paste0(if(num.folds == 1) "non ", "cross-validated")
    if (used.subset)
        attr(out, "subset") <- sub.idx
    return(out)
}

#' Predictive information criteria for Bayesian models
#'
#' Compute an efficient approximate leave-one-out cross-validation
#' using Pareto smoothed importance sampling (PSIS-LOO), or the widely
#' applicable information criterion (WAIC), also known as the Watanabe-Akaike
#' information criterion.
#'
#' @param x An object of class `hsstan`.
#' @param cores Number of cores used for parallelisation (the value of
#'        `options("mc.cores")` by default).
#' @param ... Currently ignored.
#'
#' @return
#' A `loo` object.
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' \dontshow{oldopts <- options(mc.cores=2)}
#' # continued from ?hsstan
#' loo(hs.biom)
#' waic(hs.biom)
#' \dontshow{options(oldopts)}
#'
#' @importFrom loo loo
#' @method loo hsstan
#' @aliases loo
#' @export loo
#' @export
loo.hsstan <- function(x, cores=getOption("mc.cores"), ...) {
    validate.samples(x)
    llk <- log_lik(x)
    chain.id <- rep(1:ncol(x$stanfit), each=nrow(x$stanfit))
    r.eff <- loo::relative_eff(exp(llk), chain_id=chain.id, cores=cores)
    loo <- suppressWarnings(loo::loo(llk, r_eff=r.eff, cores=cores))
    return(loo)
}

#' @rdname loo.hsstan
#' @importFrom loo waic
#' @method waic hsstan
#' @aliases waic
#' @export waic
#' @export
waic.hsstan <- function(x, cores=getOption("mc.cores"), ...) {
    validate.samples(x)
    llk <- log_lik(x)
    waic <- suppressWarnings(loo::waic(llk, cores=cores))
    return(waic)
}

#' Bayesian and LOO-adjusted R-squared
#'
#' Compute the Bayesian and the LOO-adjusted R-squared from the posterior
#' samples. For Bayesian R-squared it uses the modelled residual variance
#' (rather than the variance of the posterior distribution of the residuals).
#' The LOO-adjusted R-squared uses Pareto smoothed importance sampling LOO
#' residuals and Bayesian bootstrap.
#'
#' @param object An object of class `hsstan`.
#' @param prob Width of the posterior interval (0.95, by default). It is
#'        ignored if `summary=FALSE`.
#' @param summary Whether a summary of the distribution of the R-squared
#'        should be returned rather than the pointwise values (`TRUE` by
#'        default).
#' @param ... Currently ignored.
#'
#' @return
#' The mean, standard deviation and posterior interval of R-squared if
#' `summary=TRUE`, or a vector of R-squared values with length equal to
#' the size of the posterior sample if `summary=FALSE`.
#'
#' @references
#' Andrew Gelman, Ben Goodrich, Jonah Gabry and Aki Vehtari (2019),
#' R-squared for Bayesian regression models,
#' _The American Statistician_, 73 (3), 307-309.
#' \doi{10.1080/00031305.2018.1549100}
#'
#' Aki Vehtari, Andrew Gelman, Ben Goodrich and Jonah Gabry (2019),
#' Bayesian R2 and LOO-R2.
#' \url{https://avehtari.github.io/bayes_R2/bayes_R2.html}
#'
#' @examples
#' \dontshow{utils::example("hsstan", echo=FALSE)}
#' \dontshow{oldopts <- options(mc.cores=2)}
#' # continued from ?hsstan
#' bayes_R2(hs.biom)
#' loo_R2(hs.biom)
#' \dontshow{options(oldopts)}
#'
#' @importFrom rstantools bayes_R2
#' @method bayes_R2 hsstan
#' @aliases bayes_R2
#' @export bayes_R2
#' @export
bayes_R2.hsstan <- function(object, prob=0.95, summary=TRUE, ...) {
    validate.samples(object)
    validate.probability(prob)
    mu <- posterior_linpred(object, transform=TRUE)
    var.mu <- apply(mu, 1, stats::var)
    if (is.logistic(object))
        sigma2 <- rowMeans(mu * (1 - mu))
    else
        sigma2 <- drop(as.matrix(object$stanfit, pars="sigma"))^2
    R2 <- var.mu / (var.mu + sigma2)
    if (summary)
        R2 <- vector.summary(R2, prob)
    return(R2)
}

#' @rdname bayes_R2.hsstan
#' @importFrom rstantools loo_R2
#' @method loo_R2 hsstan
#' @aliases loo_R2
#' @export loo_R2
#' @export
loo_R2.hsstan <- function(object, prob=0.95, summary=TRUE, ...) {
    validate.samples(object)
    validate.probability(prob)
    y <- object$data[[object$model.terms$outcome]]
    mu <- posterior_linpred(object, transform=TRUE)
    ll <- log_lik(object)
    S <- nrow(mu)
    N <- ncol(mu)
    chains <- object$stanfit@sim$chains
    r.eff <- loo::relative_eff(exp(ll), chain_id=rep(1:chains, each=S / chains))
    psis <- suppressWarnings(loo::psis(-ll, r_eff=r.eff))
    mu.loo <- loo::E_loo(mu, psis, log_ratios=-ll)$value
    err.loo <- mu.loo - y

    ## set the random seed as the seed used in the first chain and ensure
    ## the old RNG state is restored on exit
    rng.state.old <- .Random.seed
    on.exit(assign(".Random.seed", rng.state.old, envir=.GlobalEnv))
    set.seed(object$stanfit@stan_args[[1]]$seed)

    ## dirichlet weights for bayesian bootstrap
    exp.draws <- matrix(stats::rexp(S * N, rate=1), nrow=S, ncol=N)
    dw <- exp.draws / rowSums(exp.draws)

    var.y <- (rowSums(sweep(dw, 2, y^2, FUN = "*")) -
              rowSums(sweep(dw, 2, y, FUN = "*"))^2) * (N / (N - 1))
    var.err.loo <- (rowSums(sweep(dw, 2, err.loo^2, FUN = "*")) -
                    rowSums(sweep(dw, 2, err.loo, FUN = "*")^2)) * (N / (N - 1))

    R2 <- 1 - var.err.loo / var.y
    R2[R2 < -1] <- -1
    R2[R2 > 1] <- 1

    if (summary)
        R2 <- vector.summary(R2, prob)
    return(R2)
}