R/brm_model.R

In brms.mmrm: Bayesian MMRMs using 'brms'

#' @title Fit an MMRM.
#' @export
#' @family models
#' @description Fit an MMRM model using `brms`.
#' @inheritSection brm_formula Parameterization
#' @return A fitted model object from `brms`, with new list elements
#'   `brms.mmrm_data` and `brms.mmrm_formula` to capture the data
#'   and formula supplied to [brm_model()]. See the explanation of the
#'   `data` argument for how the data is handled and how it relates
#'   to the data returned in the `brms.mmrm_data` attribute.
#' @param data A classed data frame from [brm_data()], or an informative
#'   prior archetype from a function like [brm_archetype_successive_cells()].
#'   Unless you supplied `model_missing_outcomes = TRUE` in [brm_formula()],
#'   [brm_model()] automatically rows with missing outcomes
#'   just prior to fitting the model with [brms::brm()].
#'   The `brms.mmrm_data` attribute
#'   in the output object is always the version of the data prior to
#'   removing these rows. See the `data` element of the returned `brms`
#'   object for the final data actually supplied to the model.
#'
#'   If you supply a non-`NULL` value for the `imputed`
#'   argument, then the `data` argument is ignored and the MMRM is fit
#'   successively to each dataset in `imputed` using [brms::brm_multiple()].
#'   Posterior draws are combined automatically for downstream post-processing
#'   unless you set `combine = FALSE` in [brm_model()].
#' @param formula An object of class `"brmsformula"` from [brm_formula()]
#'   or `brms::brmsformula()`. Should include the full mapping
#'   of the model, including fixed effects, residual correlation,
#'   and heterogeneity in the discrete-time-specific residual variance
#'   components.
#' @param ... Arguments to [brms::brm()] or [brms::brm_multiple()]
#'   other than `data`, `formula`, `prior`, and `family`.
#' @param prior Either `NULL` for default priors
#'   or a `"brmsprior"` object from `brms::prior()`.
#' @param family A `brms` family object generated by [brms::brmsfamily()].
#'   Must fit a continuous outcome variable and have the identity link.
#' @param imputed Either `NULL` (default), list of
#'   datasets generated with multiple imputation, or a `"mids"` object
#'   from the `mice` package. The `rbmi` package may offer a more appropriate
#'   method for imputation for MMRMs than `mice`. It is your responsibility
#'   to choose an imputation method appropriate for the data and model.
#'
#'   If not `NULL`, then the MMRM is fit
#'   successively to each dataset in `imputed` using [brms::brm_multiple()].
#'   Posterior draws are combined automatically for downstream post-processing
#'   unless you set `combine = FALSE` in [brm_model()], so everything at
#'   the level of [brm_marginal_draws()] will be exactly the same as
#'   a non-imputation workflow.
#'
#'   Even if you supply `imputed`, please also supply the original
#'   non-imputed dataset in the `data` argument to help with
#'   downstream post-processing.
#' @examples
#' if (identical(Sys.getenv("BRM_EXAMPLES", unset = ""), "true")) {
#' set.seed(0L)
#' data <- brm_data(
#'   data = brm_simulate_simple()$data,
#'   outcome = "response",
#'   group = "group",
#'   time = "time",
#'   patient = "patient",
#'   reference_group = "group_1",
#'   reference_time = "time_1"
#' )
#' formula <- brm_formula(
#'   data = data,
#'   baseline = FALSE,
#'   baseline_time = FALSE
#' )
#' # Optional: set the contrast option, which determines the model matrix.
#' options(contrasts = c(unordered = "contr.SAS", ordered = "contr.poly"))
#' # See the fixed effect mapping you get from the data:
#' head(brms::make_standata(formula = formula, data = data)$X)
#' # Specify a different contrast method to use an alternative
#' # mapping when fitting the model with brm_model():
#' options(
#'   contrasts = c(unordered = "contr.treatment", ordered = "contr.poly")
#' )
#' # different model matrix than before:
#' head(brms::make_standata(formula = formula, data = data)$X)
#' tmp <- utils::capture.output(
#'   suppressMessages(
#'     suppressWarnings(
#'       model <- brm_model(
#'         data = data,
#'         formula = formula,
#'         chains = 1,
#'         iter = 100,
#'         refresh = 0
#'       )
#'     )
#'   )
#' )
#' # The output is a brms model fit object with added list
#' # elements "brms.mmrm_data" and "brms.mmrm_formula" to track the dataset
#' # and formula used to fit the model.
#' model$brms.mmrm_data
#' model$brms.mmrm_formula
#' # Otherwise, the fitted model object acts exactly like a brms fitted model.
#' suppressWarnings(print(model))
#' brms::prior_summary(model)
#' }
brm_model <- function(
  data,
  formula,
  ...,
  prior = NULL,
  family = brms::brmsfamily(family = "gaussian"),
  imputed = NULL
) {
  brm_data_validate(data)
  brm_formula_validate(formula)
  brms_model_validate_family(family)
  assert(
    inherits(prior %|||% brms::prior("normal(0, 1)"), "brmsprior"),
    message = "prior arg must be a \"brmsprior\" object or NULL."
  )
  assert(
    imputed,
    is.null(.) || is.list(.) || inherits(., "mids"),
    message = paste(
      "argument 'imputed' must either be NULL, a list of data frames,",
      "or a 'mids' object from the 'mice' package. Whether in",
      "'rbmi', 'mice', or some other package, it is your responsibility",
      "to choose an imputation method appropriate for the data and model."
    )
  )
  modeled_data <- if_any(
    attr(formula, "brm_model_missing_outcomes"),
    data,
    data[!is.na(data[[attr(data, "brm_outcome")]]), ]
  )
  model <- if_any(
    is.null(imputed),
    brms::brm(
      data = modeled_data,
      formula = formula,
      prior = prior,
      ...
    ),
    brms::brm_multiple(
      data = imputed,
      formula = formula,
      prior = prior,
      ...
    )
  )
  model <- brm_model_new(model, data, formula)
  brm_model_validate(model)
  model
}

brm_model_new <- function(model, data, formula) {
  model$brms.mmrm_data <- data
  model$brms.mmrm_formula <- formula
  structure(
    model,
    class = unique(c("brms_mmrm_model", class(model)))
  )
}

brm_model_validate <- function(model) {
  assert(
    model,
    inherits(., "brms_mmrm_model"),
    inherits(., "brmsfit"),
    message = paste(
      "Please use brms.mmrm::brm_model() to fit the model.",
      "Otherwise, functions like brm_marginal_draws()",
      "in brms.mmrm may not be statistically valid."
    )
  )
  brm_data_validate(model$brms.mmrm_data)
  brm_formula_validate(model$brms.mmrm_formula)
}

brms_model_validate_family <- function(family) {
  assert(
    inherits(family, "brmsfamily"),
    message = "family must come from brms::brmsfamily()"
  )
  assert(
    identical(family$type, "real"),
    message = "family must be for continuous outcomes"
  )
  assert(
    identical(family$family, "gaussian"),
    message = "family must be gaussian"
  )
  assert(
    identical(family$link, "identity"),
    message = "family must have the identity link"
  )
  assert(
    identical(family$link_sigma, "log"),
    message = "family must link_sigma equal to 'log'"
  )
}