R/agq.R

In tulpa: Templated Unified Library for Posterior Approximation in Bayesian Hierarchical Models

#' Adaptive Gauss-Hermite quadrature for one-RE GLMMs
#'
#' @description
#' Marginal-likelihood approximation for a generalised linear mixed
#' model with one cluster-level intercept random effect. Adaptive
#' Gauss-Hermite quadrature (AGQ) generalises Laplace by replacing the
#' single-point Gaussian integral around the cluster's posterior mode
#' with `n_quad`-point quadrature. AGQ at `n_quad = 1` recovers the
#' Laplace approximation; higher `n_quad` reduces approximation error,
#' especially for clusters with few observations or non-Gaussian
#' likelihoods.
#'
#' This is an engine block, not a front door: its tuning knobs
#' (`max_iter`, `tol`, `n_quad`) sit in the signature rather than in a
#' `control` list.
#'
#' Scoped to the `lme4::glmer(..., nAGQ = N)` use case: one intercept-
#' only RE term, families `binomial`, `poisson`, or `gaussian`. For
#' multi-RE or random-slope models, use [tulpa_laplace()] (Laplace at
#' the joint mode) or HMC.
#'
#' @section Tier:
#' Tier 2 (Structured). AGQ is exact in the limit `n_quad -> infinity`
#' but at finite `n_quad` it is a controlled approximation -- same
#' epistemic class as Laplace.
#'
#' @param y Response vector.
#' @param X Fixed-effects design matrix (`n_obs x p`).
#' @param group Integer cluster labels (1-based, length `n_obs`).
#' @param n_groups Number of clusters (default `max(group)`).
#' @param family One of `"binomial"`, `"poisson"`, `"gaussian"`.
#' @param n_trials Trial sizes (binomial only; default `rep(1, n_obs)`).
#' @param sigma_eps Residual SD (gaussian only; default `1`). Held **fixed** at
#'   this value -- it is not profiled or estimated, so for a gaussian response
#'   set it to a known / pre-estimated residual SD rather than the default.
#' @param n_quad Number of Gauss-Hermite quadrature nodes per cluster.
#'   `1` recovers Laplace; common choices are `5` or `7`. Default `7`.
#' @param offset Optional fixed per-observation offset added to eta
#'   (default all zero, i.e. none).
#' @param beta_init Initial fixed-effects (default zeros).
#' @param sigma_init Initial RE SD (default `1`).
#' @param max_iter Optimiser iteration cap (default `200`).
#' @param tol Convergence tolerance (default `1e-6`).
#' @param verbose Print optimiser summary (default `FALSE`).
#'
#' @return A list with class `tulpa_fit` carrying:
#'   * `means`: named vector `c(beta, log_sigma_re)`.
#'   * `mode`: same as `means` (Gaussian fit at the optimum).
#'   * `cov`: covariance matrix of `(beta, log_sigma_re)` from inverse
#'     observed information.
#'   * `sigma_re`: estimated RE SD.
#'   * `log_marginal`: AGQ-approximated log marginal likelihood at the
#'     fitted values.
#'   * `n_quad`: quadrature order used.
#'   * `n_iter`, `converged`: optimiser diagnostics.
#'   * `inference_mode`: `"structured"`.
#'   * `inference_tier`: `2L`.
#'   * `backend`: `"agq"`.
#'
#'   An optimum at which the AGHQ objective is undefined (some group's solve
#'   fails there, so its value is the failure sentinel rather than a marginal
#'   likelihood) is an error naming the groups, not a returned fit.
#'
#' @references
#' Pinheiro, J. C., & Bates, D. M. (1995). Approximations to the
#' log-likelihood function in the nonlinear mixed-effects model.
#' *Journal of Computational and Graphical Statistics*, 4(1), 12-35.
#'
#' @seealso [tulpa_laplace()] for the joint-mode Laplace path
#'   (handles multi-RE and spatial structure).
#'
#' @examples
#' set.seed(1)
#' n_g <- 20; n_per <- 5; n <- n_g * n_per
#' group <- rep(seq_len(n_g), each = n_per)
#' x <- rnorm(n)
#' X <- cbind(1, x)
#' u <- rnorm(n_g, 0, 0.8)
#' eta <- 0.3 + 0.7 * x + u[group]
#' y <- rbinom(n, 1, plogis(eta))
#'
#' # Compare Laplace (n_quad = 1) and AGQ-7.
#' fit_lap <- agq_fit(y, X, group, family = "binomial", n_quad = 1)
#' fit_agq <- agq_fit(y, X, group, family = "binomial", n_quad = 7)
#' c(fit_lap$log_marginal, fit_agq$log_marginal)
#'
#' @export
agq_fit <- function(y, X, group,
                    n_groups = max(group),
                    family = c("binomial", "poisson", "gaussian"),
                    n_trials = NULL,
                    sigma_eps = 1.0,
                    n_quad = 7L,
                    offset = NULL,
                    beta_init = NULL,
                    sigma_init = 1.0,
                    max_iter = 200L,
                    tol = 1e-6,
                    verbose = FALSE) {

  family <- match.arg(family)
  if (!is.matrix(X)) X <- as.matrix(X)
  storage.mode(X) <- "double"
  n_obs <- length(y)
  p <- ncol(X)
  .require_fixed_effects(X, "agq_fit")
  if (nrow(X) != n_obs) {
    stop(sprintf("nrow(X) (%d) must equal length(y) (%d).",
                 nrow(X), n_obs), call. = FALSE)
  }
  if (length(group) != n_obs) {
    stop(sprintf("length(group) (%d) must equal length(y) (%d).",
                 length(group), n_obs), call. = FALSE)
  }
  group <- as.integer(group)
  if (anyNA(group)) {
    stop(sprintf(paste0(
      "`group` has missing value(s) (first at row %d). tulpa does not drop ",
      "incomplete cases; remove those rows or give them a group."),
      which(is.na(group))[1L]), call. = FALSE)
  }
  if (any(group < 1L) || any(group > n_groups)) {
    stop("`group` indices must lie in 1..n_groups.", call. = FALSE)
  }
  n_quad <- .check_n_quad(n_quad)
  # The scales reach log() and the oracle's density: a non-positive
  # `sigma_init` was "non-finite value supplied by optim", and a negative
  # `sigma_eps` was squared into a valid variance without a word
  # (gcol33/tulpa#886).
  for (arg in list(list(v = sigma_init, nm = "sigma_init"),
                   list(v = sigma_eps, nm = "sigma_eps"))) {
    if (!is.numeric(arg$v) || length(arg$v) != 1L || !is.finite(arg$v) ||
        arg$v <= 0) {
      stop(sprintf("`%s` must be a single positive number; got %s.",
                   arg$nm, paste(format(arg$v), collapse = ", ")),
           call. = FALSE)
    }
  }
  n_trials <- .normalize_n_trials(family, n_trials, n_obs)
  if (is.null(offset)) offset <- rep(0, n_obs)
  if (length(offset) != n_obs) {
    stop(sprintf("length(offset) (%d) must equal length(y) (%d).",
                 length(offset), n_obs), call. = FALSE)
  }
  # The finite guard and the support rules every other door applies: an NA
  # response and a binomial y > n_trials both surfaced as a per-group solve
  # failure advising a different beta_init / sigma_init (gcol33/tulpa#886).
  .assert_finite_model_inputs(X, y, n_trials = n_trials, offset = offset,
                              where = "agq_fit")
  .validate_family_support(family, y, n_trials = n_trials)
  if (is.null(n_trials)) n_trials <- rep(1L, n_obs)
  if (is.null(beta_init)) beta_init <- rep(0, p)
  if (length(beta_init) != p) {
    stop(sprintf("length(beta_init) (%d) must equal ncol(X) (%d).",
                 length(beta_init), p), call. = FALSE)
  }

  # Intercept-only RE: route through the shared compiled GLMM oracle, so the
  # family density (binomial / poisson / gaussian) has a single C++ source of
  # truth shared with tulpa_re_aghq() and the Gibbs sweep. The RE design is the
  # 1x1 intercept Z = 1; its covariance is one log-SD coordinate (log-Cholesky
  # of a scalar), so par = c(beta, log_sigma_re) and at n_quad = 1 this is the
  # joint Laplace. Empty groups (no obs) are handled in the oracle: they pay
  # only the N(0, sigma^2) prior. The oracle's `phi` follows the R registry
  # convention -- the residual VARIANCE for gaussian, ignored by binomial and
  # poisson -- so `sigma_eps`, which is an SD, crosses through the same named
  # bridge the SD-parameterized kernels use in the other direction.
  Z <- matrix(1, n_obs, 1L)
  orc <- cpp_glmm_oracle_make(family, .phi_to_registry(family, sigma_eps),
                              as.numeric(y),
                              as.numeric(n_trials), X, Z, group, n_groups,
                              offset = as.numeric(offset))
  par_init <- c(beta_init, log(sigma_init))

  # At n_quad > 1 supply the analytic Fisher-identity gradient (one group sweep
  # per optim step, no per-coordinate FD re-solve). At n_quad = 1 (Laplace) it
  # omits the curvature term and disagrees with the objective, so fall back to
  # optim's finite difference. The RE block is the scalar intercept (nc = 1,
  # diagonal), no theta ridge.
  if (n_quad > 1L) {
    fns <- .aghq_analytic_optim_fns(orc, 1L, FALSE, n_quad, 1.0)
    opt <- stats::optim(par = par_init, fn = fns$fn, gr = fns$gr, method = "BFGS",
                        hessian = TRUE, control = list(maxit = max_iter, reltol = tol))
  } else {
    negf <- function(par)
      -cpp_aghq_objective(par, orc, 1L, FALSE, as.integer(n_quad), 1.0)
    opt <- stats::optim(par = par_init, fn = negf, method = "BFGS", hessian = TRUE,
                        control = list(maxit = max_iter, reltol = tol))
  }

  # The objective reports a finite failure sentinel where a group's solve did
  # not succeed, and optim only moves to points it accepted -- so a run ending
  # on the sentinel started on it, and `log_marginal = -opt$value` would be that
  # sentinel while `cov` is the finite-difference curvature of the penalty
  # (gcol33/tulpa#606). There is no fit to report, so this errors rather than
  # returning one whose numbers are the starting values.
  if (.aghq_is_fail(-opt$value)) {
    gok <- .aghq_group_status(opt$par, orc, 1L, FALSE)
    stop("agq_fit: the AGHQ objective is undefined at the optimizer's own ",
         "stopping point -- the per-group solve failed for ",
         if (is.null(gok)) "at least one group"
         else .aghq_failed_group_phrase(gok),
         ". Its value is the failure sentinel, not a marginal likelihood. Try ",
         "a different `beta_init` / `sigma_init`, or a lower `n_quad`.",
         call. = FALSE)
  }

  sigma_hat <- exp(opt$par[p + 1L])

  cov_par <- tryCatch(.spd_inverse(opt$hessian),
                      error = function(e) {
                        warning("AGQ Hessian not positive definite; ",
                                "covariance set to NA.",
                                call. = FALSE)
                        matrix(NA_real_, p + 1L, p + 1L)
                      })

  par_names <- c(if (!is.null(colnames(X))) colnames(X) else paste0("beta", seq_len(p)),
                 "log_sigma_re")
  names(opt$par) <- par_names
  if (!any(is.na(cov_par))) {
    rownames(cov_par) <- colnames(cov_par) <- par_names
  }

  if (verbose) {
    cat(sprintf("agq_fit: n_quad = %d, log_marg = %.3f, converged = %s\n",
                n_quad, -opt$value, opt$convergence == 0L))
  }

  fit <- list(
    draws = matrix(NA_real_, 0L, p + 1L,
                   dimnames = list(NULL, par_names)),
    means = opt$par,
    mode = opt$par,
    cov = cov_par,
    sigma_re = sigma_hat,
    log_marginal = -opt$value,
    n_quad = as.integer(n_quad),
    n_iter = as.integer(opt$counts["function"]),
    converged = (opt$convergence == 0L),
    n_params = p + 1L,
    inference_mode = "structured",
    inference_tier = 2L,
    backend = "agq",
    param_names = par_names,
    # AGQ integrates each group's effect OUT of the marginal likelihood, so
    # there is no per-group posterior to report -- the same situation
    # tulpa_re_aghq(n_quad > 1) records. Without the reason, ranef() found no
    # `re[` columns in the 0-row draws matrix and returned an empty frame,
    # indistinguishable from a model with no random effects at all
    # (gcol33/tulpa#710).
    ranef_unavailable = paste(
      "adaptive Gauss-Hermite quadrature integrates each group's effect out of",
      "the marginal likelihood, so the fit holds sigma_re and not the per-group",
      "effects. Use mode = 're_cov_gibbs' or 're_cov_nested' for ranef()."),
    n_fixed = p
  )
  .finalize_fit(fit,
               data = list(y = y, n_trials = n_trials, model_matrix = X,
                          family = family, offset = offset,
                          phi = if (identical(family, "gaussian")) sigma_eps^2 else NULL))
}


#' Gauss-Hermite quadrature (probabilist's, exp(-z^2/2) weight)
#'
#' Computes nodes and weights for `\int f(z) (2\pi)^{-1/2} exp(-z^2/2) dz`
#' via Golub-Welsch on the Hermite Jacobi matrix. Sums of `w_k f(z_k)`
#' approximate `E_{Z ~ N(0,1)}[f(Z)]`.
#'
#' @param n Number of quadrature nodes.
#' @return List with `nodes` and `weights` (each length `n`).
#' @noRd
gauss_hermite_prob <- function(n) {
  n <- as.integer(n)
  if (n < 1L) stop("n must be >= 1.", call. = FALSE)
  if (n == 1L) return(list(nodes = 0, weights = 1))
  # Probabilist's Hermite three-term recurrence:
  #   z H_n = H_{n+1} + n H_{n-1}
  # => symmetric tridiagonal Jacobi has alpha_k = 0, beta_k = sqrt(k).
  beta <- sqrt(seq_len(n - 1L))
  Jm <- matrix(0, n, n)
  Jm[cbind(seq_len(n - 1L), seq_len(n - 1L) + 1L)] <- beta
  Jm[cbind(seq_len(n - 1L) + 1L, seq_len(n - 1L))] <- beta
  ev <- eigen(Jm, symmetric = TRUE)
  nodes <- ev$values
  # First component squared of each eigenvector x integral  weight = 1
  # (probabilist's measure normalises to 1).
  weights <- ev$vectors[1L, ]^2
  ord <- order(nodes)
  list(nodes = nodes[ord], weights = weights[ord])
}