R/predictive_checks_simulation.R

In dyadMLM: Tools for Dyadic Multilevel Models

#' Simulate response datasets for predictive checks
#'
#' `r lifecycle::badge("experimental")`
#' If a model fits, data simulated from it should resemble the observed data.
#' This function generates new response datasets for the same observations and
#' predictors.
#' Reuse them with [check_partner_dependence()], [check_dyad_residuals()], or
#' [check_dyad_outcomes()] to check whether a fitted model reproduces features of
#' the observed data. For a complete example see [check_partner_dependence()].
#'
#' @param model A fitted `glmmTMB` model.
#' @param nsim Number of complete response datasets to simulate. Default: 1000.
#' @param seed Optional seed for reproducible simulations, interpreted as an
#'   integer (see [set.seed()]). When supplied, the random-number state is restored
#'   after the function returns, or when it stops after an error.
#'
#' @return A `dyadMLM_response_simulations` object for use with
#'   [check_partner_dependence()], [check_dyad_residuals()], and [check_dyad_outcomes()].
#'
#' The result keeps all components in fitted-row order (after missing-data
#' exclusions):
#'
#' - `simulated_responses`: a matrix with `nsim` rows and one column per fitted
#'   observation. Each row is a complete simulated dataset.
#' - `observed_response` and `predicted_response`: numeric vectors with one value
#'   per fitted observation.
#' - `model_frame`: the fitted model frame (variables in the model formulas
#'   only), in the same row order.
#'
#' The `dyadMLM` attribute records the model and simulation settings, including
#' the seed.
#'
#' @section Supported models:
#' Supports unweighted `glmmTMB` models with the following families:
#' - `gaussian()`
#' - `poisson()`, `glmmTMB::compois()`, `glmmTMB::genpois()`, `glmmTMB::bell()`
#' - `glmmTMB::nbinom1()`, `glmmTMB::nbinom2()`, `glmmTMB::nbinom12()`
#' - `glmmTMB::truncated_poisson()`, `glmmTMB::truncated_nbinom1()`,
#'   `glmmTMB::truncated_nbinom2()`, `glmmTMB::truncated_compois()`,
#'   `glmmTMB::truncated_genpois()`
#' - `glmmTMB::tweedie()`
#' - `Gamma()`, `glmmTMB::ziGamma()`
#' - `glmmTMB::beta_family()`
#' - `glmmTMB::lognormal()`
#' - `glmmTMB::skewnormal()`
#' - `glmmTMB::t_family()` with more than two degrees of freedom
#' - `glmmTMB::ordinal()` (glmmTMB 1.1.15.2 or later)
#'
#' Smooth terms (`s()`) are not supported: simulations would draw new curves.
#' Fixed-effect splines such as `splines::ns()` are supported.
#'
#' Zero-inflated and hurdle versions are supported where available. Checks
#' describe the combined response, including zeros, rather than each model
#' component separately. Good agreement does not establish that the zero and
#' response components each fit well.
#'
#' Ordinal checks use category scores `1, 2, ..., K` in their fitted order,
#' matching [glmmTMB's predictions][glmmTMB::family_glmmTMB]. The checks compare
#' observed and simulated scores. The scores do not measure
#' distances on an underlying continuous scale.
#' Because each dyad contains only two people, glmmTMB tends to underestimate
#' the dyad variance in ordinal models. Simulations then show less partner
#' dependence than the data, so the partner check can flag a correct ordinal
#' model, especially with many dyads.
#'
#' The model's fitted link is used for prediction and simulation. Predictions
#' and simulated responses must be finite.
#'
#' @section Technical details:
#' Each simulation draws new random effects, if the model has any, and then new
#' responses from the fitted model. Random effects within each block are drawn
#' together using their fitted variances and correlations. This also applies to
#' random effects in the zero-inflation and dispersion models, if present.
#'
#' Fitted parameters and predictors, including any lagged responses, stay fixed.
#' The model is not refitted, and uncertainty in parameter estimates is not
#' included. This is a *plug-in predictive reference* (Gelman et al., 1996).
#' If all grouping factors are dyads or nested within dyads (such as dyad-days),
#' the simulations represent hypothetical new dyads under the same study design.
#'
#' `predicted_response` contains predicted mean responses with random effects
#' in the conditional and zero-inflation models set to zero. For zero-inflated
#' and hurdle models, this is the conditional response mean multiplied by
#' one minus the zero-component probability (Brooks et al., 2017, Appendix A;
#' \doi{10.32614/RJ-2017-066}).
#'
#' Partner-dependence checks (by default) and the response SD and largest
#' absolute deviation of [check_dyad_outcomes()] subtract these same predictions
#' from observed and simulated responses.
#' Both random effects and observation-level noise still
#' contribute to response variance. With nonlinear links, setting random effects
#' to zero generally differs from averaging predictions over them.
#'
#' @references Gelman, A., Meng, X.-L., & Stern, H. S. (1996). Posterior
#'   predictive assessment of model fitness via realized discrepancies.
#'   *Statistica Sinica, 6*(4), 733--807.
#'
#' @export
simulate_dyad_responses <- function(model, nsim = 1000, seed = NULL) {
  if (!inherits(model, "glmmTMB")) {
    stop("`model` must be a fitted `glmmTMB` model.", call. = FALSE)
  }
  if (!requireNamespace("glmmTMB", quietly = TRUE)) {
    stop("Package `glmmTMB` is required to simulate responses.", call. = FALSE)
  }
  if (!is.numeric(nsim) || !rlang::is_scalar_integerish(nsim, finite = TRUE) ||
      nsim < 1 || nsim > .Machine$integer.max) {
    stop("`nsim` must be one positive whole number.", call. = FALSE)
  }
  nsim <- as.integer(nsim)

  if (!is.null(seed)) {
    seed <- as.integer(seed)
    withr::local_seed(seed)
  }

  family <- stats::family(model)
  supported <- c(
    "gaussian", "poisson", "nbinom1", "nbinom2", "nbinom12", "compois", "genpois",
    "truncated_poisson", "truncated_nbinom1", "truncated_nbinom2",
    "truncated_compois", "truncated_genpois", "tweedie", "Gamma", "beta",
    "lognormal", "skewnormal", "bell", "t", "ordinal"
  )
  if (!family$family %in% supported) {
    stop("Unsupported family. ",
         "See the supported models in ?simulate_dyad_responses.", call. = FALSE)
  }
  if (any(stats::weights(model) != 1)) {
    stop("Predictive checks currently only support unweighted models.", call. = FALSE)
  }
  # glmmTMB fits the wiggly part of s() as a random effect, so simulate() would redraw it.
  if (any(lengths(lapply(model$modelInfo$reTrms, `[[`, "smooth_info")) > 0L)) {
    stop("Smooth terms (s()) are not supported: glmmTMB fits them as random effects, ",
         "so every simulation would draw a new curve. Use fixed-effect splines, such ",
         "as splines::ns(), instead.", call. = FALSE)
  }
  if (family$family == "t" && glmmTMB::family_params(model) <= 2) {
    stop("Student-t predictive checks require more than two degrees of freedom ",
         "so that response variances and correlations are defined.", call. = FALSE)
  }

  frame <- stats::model.frame(model)
  observed <- stats::model.response(frame) # response variable from frame (vector)
  if (family$family == "ordinal") {
    # Match predictions, which use category positions 1, 2, ..., K.
    observed <- as.numeric(observed)
    rlang::inform(paste0(
      "Ordinal categories are scored 1, 2, ..., K in both observed and ",
      "simulated data. The checks compare observed and simulated scores."
    ), .frequency = "once", .frequency_id = "dyadMLM_ordinal_scores")
  }
  if (!is.numeric(observed) || !is.null(dim(observed)) ||
      any(!is.finite(observed))) {
    stop("Expected one numeric response per fitted row.",
         call. = FALSE)
  }

  zero_inflation_model_matrix <- stats::model.matrix(model, component = "zi")
  # Since glmmTMB 1.1.15.1, zero-inflation random effects can activate the mixture.
  has_zero_inflation <- ncol(zero_inflation_model_matrix) > 0L ||
    (utils::packageVersion("glmmTMB") >= "1.1.15.1" &&
       ncol(glmmTMB::getME(model, "Zzi")) > 0L)

  conditional_design <- as.matrix(stats::model.matrix(model, component = "cond"))
  coefficient_map <- model$obj$env$map[["beta"]]
  if (!is.null(coefficient_map)) {
    # Tied coefficients move together; fixed coefficients contribute no change.
    free_groups <- outer(as.integer(coefficient_map),
                         seq_len(nlevels(coefficient_map)), "==")
    conditional_design <- conditional_design %*% (free_groups & !is.na(free_groups))
  }
  # Recognise a free overall location in any coding, including separate role means.
  # Ordinal threshold centring has not yet been validated here.
  free_conditional_intercept <- family$family != "ordinal" &&
    ncol(conditional_design) > 0L &&
    qr(conditional_design)$rank == qr(cbind(conditional_design, 1))$rank

  # newdata = NULL prevents na.exclude from padding omitted rows back in.
  predicted <- as.numeric(stats::predict(
    model, newdata = NULL, re.form = NA,
    type = if (has_zero_inflation) "conditional" else "response"
  ))

  if (has_zero_inflation) {
    # glmmTMB's re.form = NA can retain zero-inflation random effects, so use
    # the fitted design and offsets to calculate this component without them.
    zero_inflation_coefficients <- glmmTMB::fixef(model)$zi[
      colnames(zero_inflation_model_matrix)
    ]
    zero_inflation_linear_predictor <- as.numeric(
      zero_inflation_model_matrix %*% zero_inflation_coefficients
    ) + model$obj$env$data$zioffset
    # plogis(-x) calculates one minus the zero probability directly.
    predicted <- predicted * stats::plogis(-zero_inflation_linear_predictor)
  }

  # Select components whose random effects should be redrawn.
  components <- c("terms", "termszi", "termsdisp")
  # Store current simulation settings.
  original <- model$obj$env$data[components]
  if (any(vapply(original, is.null, logical(1)))) {
    stop("The fitted model has an unsupported simulation structure.", call. = FALSE)
  }
  # Restore original settings on exit, including after an error.
  on.exit(model$obj$env$data[components] <- original, add = TRUE)

  # Tell simulate() to draw new random effects for every block.
  for (component in components) {
    for (i in seq_along(original[[component]])) {
      # 2 = "random", glmmTMB's default. Forcing it also overrides any
      # set_simcodes() change by the caller that would keep random effects fixed.
      model$obj$env$data[[component]][[i]]$simCode <- 2
    }
  }


  # simulate() returns a data frame with one column per draw.
  simulated <- stats::simulate(model, nsim = nsim)
  if (family$family == "ordinal") {
    # Convert factor columns to category scores before creating the matrix.
    simulated[] <- lapply(simulated, as.numeric)
  }
  # Transpose to an nsim x nrow(frame) matrix.
  simulated <- t(as.matrix(simulated))

  if (length(predicted) != nrow(frame) || any(!is.finite(predicted)) ||
      !is.numeric(simulated) || any(!is.finite(simulated)) ||
      !identical(dim(simulated), c(nsim, nrow(frame)))) {
    stop("Expected finite predictions and one simulated response per fitted row.", call. = FALSE)
  }

  # Keep the frame with the draws so later checks resolve IDs in the same row order.
  out <- list(
    observed_response = as.numeric(observed),
    simulated_responses = simulated,
    predicted_response = predicted,
    model_frame = frame
  )

  has_random_effects <- any(lengths(lapply(model$modelInfo$reTrms, `[[`, "cnms")) > 0L)
  attr(out, "dyadMLM") <- list(
    backend = "glmmTMB",
    family = family$family,
    link = family$link,
    reference = "plug-in predictive",
    random_effects = if (has_random_effects) "new" else "no",
    parameter_uncertainty = "excluded",
    free_conditional_intercept = free_conditional_intercept,
    seed = seed
  )

  class(out) <- c("dyadMLM_response_simulations", "list")

  return(out)
}


#' Print simulated response datasets
#'
#' Prints the number and type of simulated datasets.
#'
#' @param x An object returned by [simulate_dyad_responses()].
#' @param ... Not used.
#'
#' @return `x`, invisible.
#'
#' @keywords internal
#'
#' @export
print.dyadMLM_response_simulations <- function(x, ...) {
  meta <- attr(x, "dyadMLM")
  nsim <- nrow(x$simulated_responses)
  cat("<dyadMLM response simulations>\n")
  cat(
    nsim,
    "complete", meta$family, "response",
    if (nsim == 1L) "dataset" else "datasets",
    "from", meta$backend,
    "for", length(x$observed_response), "fitted rows\n"
  )
  invisible(x)
}