R/mixstanvar.R

In RBesT: R Bayesian Evidence Synthesis Tools

#' @title Mixture distributions as \code{brms} priors
#'
#' @description Adapter function converting mixture distributions for
#'     use with \code{\link[brms]{brm}} models via the
#'     \code{\link[brms]{stanvar}} facility.
#'
#' @param ... List of mixtures to convert.
#' @param verbose Enables printing of the mixture priors when chains
#'     start to sample. Defaults to \code{FALSE}.
#'
#' @details To declare mixture priors in a \code{\link[brms]{brm}}
#'     model requires two steps: First, the mixture densities need to
#'     be converted by the adapter function \code{mixstanvar} into a
#'     \code{stanvars} object which is passed to the \code{stanvars}
#'     argument of the \code{\link[brms]{brm}} function. Doing so
#'     extends the Stan code and data generated by
#'     \code{\link[brms]{brm}} such that the mixture densities can be
#'     used as priors within the \code{\link[brms]{brm}} model. The
#'     second step is to assign parameters of the
#'     \code{\link[brms]{brm}} model to the mixture density as prior
#'     using the \code{\link[brms]{set_prior}} command of \code{brms}.
#'
#'  The adapter function translates the mixture distributions as
#'  defined in \code{R} to the respective mixture distribution in
#'  Stan. Within Stan the mixture distributions are named in
#'  accordance to the \code{R} functions used to create the respective
#'  mixture distributions. That is, a mixture density of normals
#'  created by \code{\link{mixnorm}} is referred to as
#'  \code{mixnorm_lpdf} in Stan such that one can refer to the density
#'  as \code{mixnorm} within the \code{\link[brms]{set_prior}}
#'  functions (the suffix \code{_lpdf} is automatically added by
#'  \code{\link[brms]{brm}}). The arguments to these mixture
#'  distributions depend on the specific distribution type as follows:
#'
#' \tabular{cl}{
#' Density \tab Arguments \cr
#' \code{mixbeta(w, a, b)} \tab \code{w} weights, \code{a} shapes, \code{b} shapes \cr
#' \code{mixgamma(w, a, b)} \tab \code{w} weights, \code{a} shapes, \code{b} inverse scales \cr
#' \code{mixnorm(w, m, s)} \tab \code{w} weights, \code{m} means, \code{s} standard deviations \cr
#' \code{mixmvnorm(w, m, sigma_L)} \tab \code{w} weights, \code{m} means, \code{sigma_L} cholesky factors of covariances \cr
#' }
#'
#' These arguments to the mixture densities refer to the different
#' density parameters and are automatically extracted from the
#' mixtures when converted. Important here are the argument names as
#' these must be used to declare the mixture prior. For each mixture
#' to convert as part of the \code{...}  argument to \code{mixstanvar}
#' a label is determined using the name given in the list. In case no
#' name is given, then the name of the \code{R} object itself is
#' used. To declare a prior for a parameter the mixture distribution
#' must be used with it's arguments following the convention
#' \code{label_argument}. Please refer to the examples section for an
#' illustration.
#'
#' \strong{Note:} Models created by \code{\link[brms]{brm}} do use by
#' default a data-dependent centering of covariates leading to a shift
#' of the overall intercept. This is commonly not desirable in
#' applications of this functionality. It is therefore strongly
#' recommended to pass the option \code{center=FALSE} as argument to
#' the \code{brms} formula created with the \code{\link[brms]{bf}}
#' function as demonstrated with the example below.
#'
#' @return \code{stanvars} object to be used as argument to the
#'     \code{stanvars} argument of a \code{\link[brms]{brm}} model.
#'
#' @examples
#' \dontrun{
#' # The mixstanvar adapter requires the optional packages brms and glue
#' stopifnot(require("brms"), require("glue"))
#' 
#' # Assume we prefer a logistic regression MCMC analysis rather than a
#' # beta-binomial analysis for the responder endpoint of the ankylosing
#' # spondylitis (AS) example. Reasons to prefer a regression analysis is
#' # to allow for baseline covariate adjustments, for example.
#' map_AS_beta <- mixbeta(c(0.62, 19.2, 57.8), c(0.38, 3.5, 9.4))
#' 
#' # First we need to convert the beta mixture to a respective mixture on
#' # the log odds scale and approximate it with a normal mixture density.
#' map_AS_samp <- rmix(map_AS_beta, 1E4)
#' map_AS <- mixfit(logit(map_AS_samp), type="norm", Nc=2)
#' 
#' # Trial results for placebo and secukinumab.
#' trial <- data.frame(n=c(6, 24),
#'                     r=c(1, 15),
#'                     arm=factor(c("placebo", "secukinumab")))
#' 
#' # Define brms model such that the overall intercept corresponds to the
#' # placebo response rate on the logit scale. NOTE: The use of
#' # center=FALSE is required here as detailed in the note above.
#' model <- bf(r | trials(n) ~ 1 + arm, family=binomial, center=FALSE)
#' # to obtain detailed information on the declared model parameters use
#' # get_prior(model, data=trial)
#' # declare model prior with reference to mixture normal map prior...
#' model_prior <- prior(mixnorm(map_w, map_m, map_s), coef=Intercept) +
#'     prior(normal(0, 2), class=b)
#'
#' # ... which must be made available to brms using the mixstanvar adapter.
#' # Note that the map_AS prior is labeled "map" as referred to in the
#' # previous prior declaration.
#' analysis <- brm(model, data=trial, prior=model_prior,
#'                 stanvars=mixstanvar(map=map_AS),
#'                 seed=365634, refresh=0)
#' 
#' # Let's compare the logistic regression estimate for the probability
#' # of a positive treatment effect (secukinumab response rate exceeding
#' # the response rate of placebo) to the direct beta-binomial analysis:
#' hypothesis(analysis, "armsecukinumab > 0")
#' 
#' post_secukinumab <- postmix(mixbeta(c(1, 0.5, 1)), r=15, n=24)
#' post_placebo     <- postmix(map_AS_beta, r=1, n=6)
#' pmixdiff(post_secukinumab,  post_placebo, 0, lower.tail=FALSE)
#' # The posterior probability for a positive treatment effect
#' # is very close to unity in both cases.
#' }
#' @export
mixstanvar <- function(..., verbose=FALSE) {
    .assert_namespace("brms")
    .assert_namespace("glue")
    default_variable_names <- lapply(rlang::enquos(...), rlang::as_label)
    mixpriors <- list(...)
    if(is.null(names(mixpriors))) {
        variable_names <- default_variable_names
    } else {
        variable_names <- names(mixpriors)
        not_set <- which(variable_names == "")
        variable_names[not_set] <- default_variable_names[not_set]
    }
    assert_list(variable_names, "character", unique=TRUE)
    sv <- mix2brms(mixpriors[[1]], variable_names[[1]], verbose)
    for(i in seq_len(length(mixpriors)-1)) {
        mix <- mixpriors[[i+1]]
        variable <- variable_names[[i+1]]
        sv <- sv + mix2brms(mix, variable, verbose)
    }
    includes_density <- function(density) any(sapply(mixpriors, inherits, density))

    if(includes_density("mvnormMix")) {
        sv <- sv + brms::stanvar(name="mixmvnorm_lpdf", scode="
real mixmvnorm_lpdf(vector y, vector w, vector[] m, matrix[] L) {
  int Nc = rows(w);
  vector[Nc] lp_mix;
  for(i in 1:Nc) {
     lp_mix[i] = multi_normal_cholesky_lpdf(y | m[i], L[i]);
  }
  return log_sum_exp(log(w) + lp_mix);
}", block="functions")
    }
    mixdensity_template <- function(mixdens, standens, arg1, arg2) {
        brms::stanvar(name=glue::glue("{mixdens}_lpdf"), scode=glue::glue("
real {{mixdens}}_lpdf(real y, vector w, vector {{arg1}}, vector {{arg2}}) {
  int Nc = rows(w);
  vector[Nc] lp_mix;
  for(i in 1:Nc) {
     lp_mix[i] = {{standens}}_lpdf(y | {{arg1}}[i], {{arg2}}[i]);
  }
  return log_sum_exp(log(w) + lp_mix);
}", .open="{{", .close="}}"), block="functions")
    }
    if(includes_density("normMix")) {
        sv <- sv + mixdensity_template("mixnorm", "normal", "m", "s")
    }
    if(includes_density("betaMix")) {
        sv <- sv + mixdensity_template("mixbeta", "beta", "a", "b")
    }
    if(includes_density("gammaMix")) {
        sv <- sv + mixdensity_template("mixgamma", "gamma", "a", "b")
    }
    
    sv
}

#' @keywords internal
mix2brms <- function(mix, name, verbose=FALSE) UseMethod("mix2brms")

#' @keywords internal
mix2brms.default <- function(mix, name, verbose=FALSE) {
    stop("Mixture density not supported in mixstanvar.")
}

#' @keywords internal
mix2brms.mvnormMix <- function(mix, name, verbose=FALSE) {
    Nc <- ncol(mix)
    p <- mvnormdim(mix[-1,1])
    Sigma <- array(NA, dim=c(Nc, p, p))
    for(i in 1:Nc) {
        Rho_c <- diag(nrow=p)
        Rho_c[lower.tri(Rho_c)] <- mix[(1+2*p+1):nrow(mix),i,drop=FALSE]
        Rho_c[upper.tri(Rho_c)] <- t(Rho_c)[upper.tri(Rho_c)]
        s <- mix[(1+p+1):(1+p+p),i,drop=TRUE]
        Sigma[i,,] <- diag(s, nrow=p) %*% Rho_c %*% diag(s, nrow=p)
    }
    prefix <- paste0(name, "_")
    mvprior <-  brms::stanvar(Nc, glue::glue("{prefix}Nc")) +
        brms::stanvar(p, glue::glue("{prefix}p")) +
        brms::stanvar(array(mix[1,,drop=TRUE], dim=Nc), glue::glue("{prefix}w"), scode=glue::glue("vector[{prefix}Nc] {prefix}w;")) +
        brms::stanvar(t(mix[2:(p+1),,drop=FALSE]), glue::glue("{prefix}m"), scode=glue::glue("vector[{prefix}p] {prefix}m[{prefix}Nc];")) +
        brms::stanvar(Sigma, glue::glue("{prefix}sigma"), scode=glue::glue("matrix[{prefix}p, {prefix}p] {prefix}sigma[{prefix}Nc];")) +
        brms::stanvar(scode=glue::glue("
matrix[{{prefix}}p, {{prefix}}p] {{prefix}}sigma_L[{{prefix}}Nc];
for (i in 1:{{prefix}}Nc) {
    {{prefix}}sigma_L[i] = cholesky_decompose({{prefix}}sigma[i]);
}", .open="{{", .close="}}"), block="tdata")
    if(verbose) {
        mvprior <- mvprior +
            brms::stanvar(scode=glue::glue('
print("Mixture prior {{name}}");
for(i in 1:{{prefix}}Nc) {
  print("Component ", i, ": w     = ", {{prefix}}w[i]);
  print("Component ", i, ": m     = ", {{prefix}}m[i]);
  print("Component ", i, ": Sigma = ", {{prefix}}sigma[i]);
}
', .open="{{", .close="}}"), position="end", block="tdata")
    }
    mvprior
}

#' @keywords internal
mix2brms.normMix <- function(mix, name, verbose=FALSE) {
    .declare_scalar_mixture_components(mix, c("w", "m", "s"), name, verbose)
}

#' @keywords internal
mix2brms.betaMix <- function(mix, name, verbose=FALSE) {
    .declare_scalar_mixture_components(mix, c("w", "a", "b"), name, verbose)
}

#' @keywords internal
mix2brms.gammaMix <- function(mix, name, verbose=FALSE) {
    .declare_scalar_mixture_components(mix, c("w", "a", "b"), name, verbose)
}

#' @keywords internal
.declare_scalar_mixture_components <- function(mix, vars, name, verbose=FALSE) {
    assert_that(length(vars) == 3, msg="Each mixture component density is expected to have 3 arguments.")
    Nc <- ncol(mix)
    prefix <- paste0(name, "_")
    sv <-  brms::stanvar(Nc, glue::glue("{prefix}Nc"))
    for(i in 1:3) {
        arg <- vars[i]
        sv <- sv + brms::stanvar(array(mix[i,,drop=TRUE], dim=Nc), glue::glue("{prefix}{arg}"), scode=glue::glue("vector[{prefix}Nc] {prefix}{arg};"))
    }
    if(verbose) {
        sv <- sv +
            brms::stanvar(name=paste0("verbose_", name), scode=glue::glue('
print("Mixture prior {{name}}");
for(i in 1:{{prefix}}Nc) {
  print("Component ", i, ": w = ", {{prefix}}{vars[1]}[i]);
  print("Component ", i, ": m = ", {{prefix}}{vars[2]}[i]);
  print("Component ", i, ": s = ", {{prefix}}{vars[3]}[i]);
}
', .open="{{", .close="}}"), position="end", block="tdata")
    }
    sv
}

#' @keywords internal
.assert_namespace <- function(package) {
    assert_that(requireNamespace(package, quietly=TRUE),
                msg=paste0("Package ", package, " must be installed first."))
}