R/ordinal.R

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

# ordinal.R
# ------------------------------------------------------------------------------
# Ordinal (ordered K-class) cumulative-link regression via a Laplace fit. For
# an ordered response the cumulative probabilities are
#   P(y <= j | x) = F(c_j - x'beta),   j = 1..K-1,
# with F the logistic CDF (proportional odds, the default) or the standard
# normal CDF (cumulative probit), ordered cutpoints c_1 < ... < c_{K-1}, and NO
# fixed-effect intercept (the cutpoints absorb the level). The penalized mode
# is found with L-BFGS on the negative log-posterior (a ridge prior on beta and
# the cutpoints) and summarized by a Laplace approximation from the numerical
# Hessian at the mode. Ordering is guaranteed by the reparameterization
# c_1 = g_1, c_k = c_{k-1} + exp(g_k).
# ------------------------------------------------------------------------------

# The cumulative link CDF of a link name.
.ordinal_link_cdf <- function(link) {
  if (identical(link, "probit")) stats::pnorm else stats::plogis
}

# Class probabilities [n x K] of the cumulative-link model at location `eta`
# and ordered cutpoints `cuts`: P(y = j) = F(c_j - eta) - F(c_{j-1} - eta),
# with c_0 = -Inf and c_K = Inf.
.ordinal_class_probs <- function(eta, cuts, pfun) {
  Fmat <- pfun(outer(-eta, cuts, "+"))            # [n x K1], F_{ij} = F(c_j - eta_i)
  cbind(Fmat, 1) - cbind(0, Fmat)
}

# Negative log-posterior in the unconstrained (beta, gamma) parameterization.
# `pfun` is the cumulative link CDF (plogis / pnorm).
.ordinal_nll <- function(par, X, cls, K, tau_b, tau_c, pfun) {
  p  <- ncol(X); K1 <- K - 1L
  beta <- par[seq_len(p)]
  gam  <- par[(p + 1L):(p + K1)]
  cuts <- cumsum(c(gam[1L], exp(gam[-1L])))       # ordered cutpoints
  eta  <- as.numeric(X %*% beta)
  n    <- length(eta)
  pobs <- .ordinal_class_probs(eta, cuts, pfun)[cbind(seq_len(n), cls)]
  pobs <- pmax(pobs, 1e-12)
  -(sum(log(pobs)) - 0.5 * tau_b * sum(beta^2) - 0.5 * tau_c * sum(gam^2))
}

#' Ordinal (ordered K-class) cumulative-logit regression via Laplace
#'
#' @description
#' Fits a proportional-odds cumulative-logit model for an ordered factor
#' response: `P(y <= j | x) = plogis(c_j - x'beta)` with ordered cutpoints
#' `c_1 < ... < c_{K-1}`. The penalized mode (a ridge prior on the coefficients
#' and cutpoints) is found by L-BFGS and summarized by a Laplace approximation.
#' There is no separate intercept -- the cutpoints carry the baseline levels.
#'
#' @param formula Model formula; the response must be an ordered (or coercible)
#'   factor with >= 3 levels, each observed at least once. An intercept in
#'   `formula` is dropped; `y ~ 1` fits the cutpoints alone. Incomplete rows
#'   are refused, not dropped.
#' @param data A data frame.
#' @param link Cumulative link: `"logit"` (proportional odds, default) or
#'   `"probit"`.
#' @param beta_prior Fixed-effect prior as `list(mean, sd)`: a mean-zero
#'   (`mean = 0`) Gaussian ridge on every coefficient with scalar SD `sd`
#'   (default the engine default, `prior_normal(0, 2.5)`).
#' @param cut_prior_sd SD of the mean-zero Gaussian ridge prior on the cutpoint
#'   parameters (default 10).
#' @param control List of numerical knobs: `max_iter` (default 200), `n_draws`
#'   (default 2000), `seed`.
#'
#' @return A `tulpa_fit` (subclass `tulpa_ordinal`) with `coefficients`
#'   (covariate effects) and `cutpoints` at the posterior mode, `draws` (the
#'   Laplace Gaussian mapped to ordered cutpoints), `log_marginal`, `levels`.
#'   The draws are the posterior [coef()], [vcov()], [summary.tulpa_fit()] and
#'   [confint()] report; [coef()] omits the cutpoints, as `MASS::polr` does.
#'
#' @seealso [tulpa_multinomial()] for the nominal (unordered) case.
#' @examples
#' \donttest{
#' set.seed(1)
#' n <- 400L; x <- rnorm(n)
#' cuts <- c(-1, 0.5, 2); eta <- 0.8 * x
#' Fm <- plogis(outer(-eta, cuts, "+")); P <- cbind(Fm, 1) - cbind(0, Fm)
#' y <- ordered(apply(P, 1, function(pr) sample.int(4L, 1L, prob = pr)))
#' fit <- tulpa_ordinal(y ~ x, data = data.frame(y = y, x = x))
#' fit$coefficients; fit$cutpoints
#' }
#' @export
tulpa_ordinal <- function(formula, data, link = c("logit", "probit"),
                          beta_prior = .tulpa_default_beta_prior("ordinal"),
                          cut_prior_sd = 10,
                          control = list()) {
  link <- match.arg(link)
  tulpa_check_control(control, .CONTROL_KEYS$ordinal, "tulpa_ordinal")
  beta_prior_sd <- .beta_prior_ridge_sd(beta_prior, .tulpa_prior_sd("ordinal"))
  pfun <- .ordinal_link_cdf(link)
  qfun <- if (link == "probit") stats::qnorm else stats::qlogis
  max_iter <- as.integer(control$max_iter %||% 200L)
  n_draws  <- as.integer(control$n_draws %||% 2000L)

  cm <- .categorical_model_frame(formula, data, "tulpa_ordinal", ordered = TRUE)
  y  <- cm$y
  K  <- nlevels(y)
  if (K < 3L) {
    stop("tulpa_ordinal() needs a response with >= 3 ordered levels; for 2 use ",
         "family = \"binomial\".", call. = FALSE)
  }
  K1  <- K - 1L
  cls <- as.integer(y)
  X   <- cm$X
  X   <- X[, setdiff(colnames(X), "(Intercept)"), drop = FALSE]   # cutpoints carry it
  # `y ~ 1` leaves no covariate: the cutpoints alone are the model (the
  # marginal class frequencies), which MASS::polr fits too (gcol33/tulpa#886).
  p   <- ncol(X)

  tau_b <- 1 / beta_prior_sd^2
  tau_c <- 1 / cut_prior_sd^2

  # Init: zero effects, evenly spread cutpoints on the link scale.
  c0   <- qfun(seq_len(K1) / K)
  par0 <- c(rep(0, p), c(c0[1L], log(pmax(diff(c0), 1e-3))))

  opt <- stats::optim(par0, .ordinal_nll, method = "BFGS", hessian = TRUE,
                      control = list(maxit = max_iter),
                      X = X, cls = cls, K = K, tau_b = tau_b, tau_c = tau_c,
                      pfun = pfun)

  par  <- opt$par
  beta <- par[seq_len(p)]
  gam  <- par[(p + 1L):(p + K1)]
  cuts <- cumsum(c(gam[1L], exp(gam[-1L])))
  names(beta) <- colnames(X)
  names(cuts) <- paste0(levels(y)[-K], "|", levels(y)[-1L])

  # Laplace marginal from the numerical Hessian of the negative log-joint.
  H <- opt$hessian
  ch <- tryCatch(chol(H), error = function(e) chol(H + diag(1e-6, nrow(H))))
  logdetH <- 2 * sum(log(diag(ch)))
  dim_b <- length(par)
  log_marginal <- -opt$value + 0.5 * dim_b * log(2 * pi) - 0.5 * logdetH

  V <- chol2inv(ch)
  pn <- c(colnames(X), names(cuts))
  dimnames(V) <- list(pn, pn)
  .seed_scoped(control$seed)
  draws <- .ps_rmvnorm(n_draws, setNames(par, pn), V)

  # The optimizer works in the unconstrained (c1, log-increment) space; map the
  # cutpoint block of every draw to ordered cutpoints so the draws the
  # coefficient methods read are labelled in cutpoint space. The map is
  # nonlinear, so the cutpoint posterior is not Gaussian and the draws, not a
  # closed-form Gaussian, are what this fit reports.
  cut_cols <- (p + 1L):(p + K1)
  draws[, cut_cols] <- t(apply(draws[, cut_cols, drop = FALSE], 1L,
                               function(g) cumsum(c(g[1L], exp(g[-1L])))))
  means <- setNames(c(beta, cuts), pn)

  fit <- list(
    coefficients = beta, cutpoints = cuts, draws = draws,
    means = means, param_names = pn,
    log_marginal = log_marginal, converged = opt$convergence == 0,
    levels = levels(y), n_classes = K, link = link,
    family = if (link == "probit") "ordinal_probit" else "ordinal",
    formula = formula, model_matrix = X, y = y, N = length(y),
    backend = "ordinal_laplace",
    inference_tier = 2L, inference_mode = "structured", draws_kind = "iid"
  )
  class(fit) <- c("tulpa_ordinal", "tulpa_categorical", "tulpa_fit")
  fit
}

# The covariate effects of the fixed-effect table, cutpoints omitted as
# MASS::polr's coef() omits its zeta; vcov() and summary() keep the cutpoints.
#' @rdname tulpa_ordinal
#' @param object A `tulpa_ordinal` fit.
#' @param ... Ignored.
#' @return `coef()` returns the covariate effects without the cutpoints, as
#'   `MASS::polr()` omits its `zeta`; `vcov()` and `summary()` keep them.
#' @export
coef.tulpa_ordinal <- function(object, ...) {
  .ordinal_table_split(object)$coefficients
}

# The fixed-effect table's estimates split into covariate effects and
# cutpoints, so coef() and print() read the posterior summary() reports.
#' @keywords internal
.ordinal_table_split <- function(object) {
  tab <- .fit_fixed_table(object)
  est <- stats::setNames(tab$estimate, tab$term)
  cut <- names(object$cutpoints)
  list(coefficients = est[setdiff(names(est), cut)], cutpoints = est[cut])
}

#' @rdname tulpa_ordinal
#' @param x A `tulpa_ordinal` fit.
#' @export
print.tulpa_ordinal <- function(x, ...) {
  cat(sprintf("Ordinal cumulative %s (%d ordered levels), Laplace fit\n",
              x$link %||% "logit", x$n_classes))
  cat(sprintf("log marginal: %.2f\n\nCoefficients:\n", x$log_marginal))
  sp <- .ordinal_table_split(x)
  print(round(sp$coefficients, 4))
  cat("\nCutpoints:\n"); print(round(sp$cutpoints, 4))
  invisible(x)
}