R/convergence.R

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

# =============================================================================
# convergence.R -- MCMC convergence diagnostics (generic, model-agnostic).
#
# Rank-normalized split-Rhat, folded split-Rhat, bulk / tail / mean / sd /
# quantile effective sample size, and Monte Carlo standard errors, following
# Vehtari, Gelman, Simpson, Carpenter & Burkner (2021), "Rank-normalization,
# folding, and localization: an improved Rhat for assessing convergence of
# MCMC", Bayesian Analysis 16(2):667-718. Native implementation -- no
# posterior / coda dependency.
#
# The scalar estimators here reproduce posterior::rhat / ess_* / mcse_* to
# machine precision on a fixed draws array (see tests/testthat/test-convergence.R).
#
# These estimators are the single source of truth for Rhat / ESS / MCSE in the
# tulpa ecosystem: the plotting layer (plot_rhat, plot_ess, diagnostic_summary,
# check_diagnostics) and downstream model packages (tulpaObs, tulpaRatio) all
# reach them through `diagnostics()` (diagnostics.R). They are location/scale-
# invariant per parameter, so they may be computed on engine-scale draws.
# =============================================================================

# --- Scalar estimators on a prepared [N_iter x M_chain] matrix --------------

# Autocovariance at lags 0..n-1 for one chain, via FFT (biased estimator,
# divisor n so element 1 is the population variance).
.tulpa_autocov <- function(x) {
  n <- length(x)
  x <- x - mean(x)
  nfft <- 2^ceiling(log2(2 * n))                 # zero-pad to avoid wraparound
  f <- stats::fft(c(x, rep(0, nfft - n)))
  acov <- Re(stats::fft(f * Conj(f), inverse = TRUE)) / nfft
  acov[seq_len(n)] / n
}

# Split each chain in half: [N x M] -> [floor(N/2) x 2M], so a single long
# chain still yields a within/between comparison.
.tulpa_split_chains <- function(sims) {
  n <- nrow(sims)
  if (n < 2L) return(sims)
  half <- n %/% 2L
  cbind(sims[seq_len(half), , drop = FALSE],
        sims[(n - half + 1L):n, , drop = FALSE])
}

# Rank-normalize the pooled draws (average ranks for ties, then the normal
# quantile of the Blom-transformed fractional rank, c = 3/8 so the divisor is
# S + 1/4 -- matches posterior::z_scale).
.tulpa_rank_normalize <- function(sims) {
  r <- rank(c(sims), ties.method = "average")
  s <- length(r)
  z <- stats::qnorm((r - 3 / 8) / (s + 1 / 4))
  matrix(z, nrow = nrow(sims), ncol = ncol(sims))
}

# Fold draws about their pooled median (detects scale / tail divergence).
.tulpa_fold <- function(sims) abs(sims - stats::median(c(sims)))

# Classic Gelman-Rubin potential scale reduction on a prepared (split,
# rank-normalized) [N x M] matrix.
.tulpa_rhat_basic <- function(sims) {
  n <- nrow(sims); m <- ncol(sims)
  if (n < 2L || m < 2L) return(NA_real_)
  chain_mean <- colMeans(sims)
  chain_var  <- apply(sims, 2L, stats::var)
  W <- mean(chain_var)
  if (!is.finite(W) || W <= 0) return(NA_real_)
  B <- n * stats::var(chain_mean)                # var_between
  sqrt((B / W + n - 1) / n)
}

# Bulk Rhat: rank-normalized split-Rhat.
.tulpa_rhat_bulk <- function(sims) {
  .tulpa_rhat_basic(.tulpa_rank_normalize(.tulpa_split_chains(sims)))
}

# Folded Rhat: rank-normalized split-Rhat of the folded draws.
.tulpa_rhat_fold <- function(sims) {
  .tulpa_rhat_basic(.tulpa_rank_normalize(.tulpa_split_chains(.tulpa_fold(sims))))
}

# Improved Rhat: the larger of the bulk and folded values (Vehtari et al. 2021).
.tulpa_rhat <- function(sims) {
  rb <- .tulpa_rhat_bulk(sims)
  rf <- .tulpa_rhat_fold(sims)
  if (is.na(rb) && is.na(rf)) return(NA_real_)
  max(rb, rf, na.rm = TRUE)
}

# Effective sample size for a [N x M] matrix, via the multi-chain
# autocorrelation estimator and Geyer's initial monotone-sequence truncation
# (the estimator used by Stan and posterior::ess_rfun).
.tulpa_ess <- function(sims) {
  if (is.null(dim(sims))) sims <- matrix(sims, ncol = 1L)
  n <- nrow(sims); m <- ncol(sims)
  if (n < 4L) return(NA_real_)
  acov <- vapply(seq_len(m), function(j) .tulpa_autocov(sims[, j]), numeric(n))
  chain_mean <- colMeans(sims)
  mean_var <- mean(acov[1L, ]) * n / (n - 1)
  if (!is.finite(mean_var) || mean_var <= 0) return(NA_real_)
  var_plus <- mean_var * (n - 1) / n
  if (m > 1L) var_plus <- var_plus + stats::var(chain_mean)
  if (!is.finite(var_plus) || var_plus <= 0) return(NA_real_)

  rho_hat_t <- numeric(n)
  rho_hat_t[1L] <- 1                                   # lag 0
  rho_hat_odd <- 1 - (mean_var - mean(acov[2L, ])) / var_plus
  rho_hat_t[2L] <- rho_hat_odd
  rho_hat_even <- 1
  t <- 0L
  while (t < n - 5L && is.finite(rho_hat_even + rho_hat_odd) &&
         (rho_hat_even + rho_hat_odd) > 0) {
    t <- t + 2L
    rho_hat_even <- 1 - (mean_var - mean(acov[t + 1L, ])) / var_plus
    rho_hat_odd  <- 1 - (mean_var - mean(acov[t + 2L, ])) / var_plus
    if ((rho_hat_even + rho_hat_odd) >= 0) {
      rho_hat_t[t + 1L] <- rho_hat_even
      rho_hat_t[t + 2L] <- rho_hat_odd
    }
  }
  max_t <- t
  if (rho_hat_even > 0) rho_hat_t[max_t + 1L] <- rho_hat_even
  # initial monotone sequence: pair sums must be non-increasing
  k <- 0L
  while (k <= max_t - 4L) {
    k <- k + 2L
    if (rho_hat_t[k + 1L] + rho_hat_t[k + 2L] >
        rho_hat_t[k - 1L] + rho_hat_t[k]) {
      rho_hat_t[k + 1L] <- (rho_hat_t[k - 1L] + rho_hat_t[k]) / 2
      rho_hat_t[k + 2L] <- rho_hat_t[k + 1L]
    }
  }
  tau_hat <- -1 + 2 * sum(rho_hat_t[seq_len(max_t)]) + rho_hat_t[max_t + 1L]
  ess <- n * m
  tau_bound <- 1 / log10(ess)                          # cap unstable estimates
  if (!is.finite(tau_hat) || tau_hat < tau_bound) tau_hat <- tau_bound
  ess / tau_hat
}

# Bulk ESS: ESS of the rank-normalized split draws.
.tulpa_ess_bulk <- function(sims) {
  .tulpa_ess(.tulpa_rank_normalize(.tulpa_split_chains(sims)))
}

# ESS for a quantile probability: ESS of the split tail-indicator draws.
.tulpa_ess_quantile <- function(sims, prob) {
  if (prob == 1) {
    len <- length(sims)
    prob <- (len - 0.5) / len
  }
  q <- stats::quantile(c(sims), prob, names = FALSE, na.rm = TRUE)
  ind <- matrix(as.numeric(sims <= q), nrow(sims))
  .tulpa_ess(.tulpa_split_chains(ind))
}

# Tail ESS: the smaller of the 5% and 95% tail-quantile ESS values.
.tulpa_ess_tail <- function(sims) {
  min(.tulpa_ess_quantile(sims, 0.05), .tulpa_ess_quantile(sims, 0.95))
}

# ESS for the mean: ESS of the split (non-rank-normalized) draws.
.tulpa_ess_mean <- function(sims) {
  .tulpa_ess(.tulpa_split_chains(sims))
}

# ESS for the standard deviation: ESS of the split squared centred draws.
.tulpa_ess_sd <- function(sims) {
  sc <- sims - mean(c(sims))
  .tulpa_ess(.tulpa_split_chains(sc^2))
}

# Monte Carlo standard error of the posterior mean.
.tulpa_mcse_mean <- function(sims) {
  stats::sd(c(sims)) / sqrt(.tulpa_ess_mean(sims))
}

# Monte Carlo standard error of the posterior standard deviation.
.tulpa_mcse_sd <- function(sims) {
  sc <- sims - mean(c(sims))
  ess <- .tulpa_ess_sd(sims)                # same split-squared-centred ESS
  e_var <- mean(sc^2)
  var_var <- (mean(sc^4) - e_var^2) / ess
  sqrt(var_var / e_var / 4)
}

# Monte Carlo standard error of a posterior quantile (Vehtari et al. 2021,
# beta-quantile interval around the tail-indicator ESS).
.tulpa_mcse_quantile <- function(sims, prob) {
  ess <- .tulpa_ess_quantile(sims, prob)
  p <- c(0.1586553, 0.8413447)                         # +- 1 sd of N(0, 1)
  a <- stats::qbeta(p, ess * prob + 1, ess * (1 - prob) + 1)
  ssims <- sort(c(sims))
  s <- length(ssims)
  th1 <- ssims[max(floor(a[1L] * s), 1L)]
  th2 <- ssims[min(ceiling(a[2L] * s), s)]
  (th2 - th1) / 2
}

# --- Measure registry -------------------------------------------------------
# Each measure maps a prepared [N_iter x M_chain] matrix (and quantile probs)
# to a named numeric vector. Adding a column is O(1): register one entry.
# `probs`-driven measures expand to one column per probability.

.tulpa_diag_measures <- list(
  rhat      = function(sims, probs) c(rhat = .tulpa_rhat(sims)),
  rhat_bulk = function(sims, probs) c(rhat_bulk = .tulpa_rhat_bulk(sims)),
  rhat_fold = function(sims, probs) c(rhat_fold = .tulpa_rhat_fold(sims)),
  ess_bulk  = function(sims, probs) c(ess_bulk = .tulpa_ess_bulk(sims)),
  ess_tail  = function(sims, probs) c(ess_tail = .tulpa_ess_tail(sims)),
  ess_mean  = function(sims, probs) c(ess_mean = .tulpa_ess_mean(sims)),
  ess_sd    = function(sims, probs) c(ess_sd = .tulpa_ess_sd(sims)),
  mcse_mean = function(sims, probs) c(mcse_mean = .tulpa_mcse_mean(sims)),
  mcse_sd   = function(sims, probs) c(mcse_sd = .tulpa_mcse_sd(sims)),
  ess_quantile = function(sims, probs)
    stats::setNames(vapply(probs, function(p) .tulpa_ess_quantile(sims, p),
                           numeric(1)), paste0("ess_q", probs * 100)),
  mcse_quantile = function(sims, probs)
    stats::setNames(vapply(probs, function(p) .tulpa_mcse_quantile(sims, p),
                           numeric(1)), paste0("mcse_q", probs * 100))
)

# Column names a measure set produces, in order (deterministic given probs).
.tulpa_measure_cols <- function(measures, probs) {
  unlist(lapply(measures, function(mn) {
    if (mn == "ess_quantile")  return(paste0("ess_q",  probs * 100))
    if (mn == "mcse_quantile") return(paste0("mcse_q", probs * 100))
    mn
  }), use.names = FALSE)
}

# Evaluate a measure set on one parameter's draws -> named numeric vector.
.tulpa_measure_eval <- function(sims, measures, probs) {
  unlist(lapply(measures, function(mn) .tulpa_diag_measures[[mn]](sims, probs)))
}

# --- Draws provenance gate --------------------------------------------------
# A fit's `$draws` may hold an autocorrelated MCMC chain or an i.i.d. sample
# from a deterministic approximation (a nested-Laplace / VI / SMC fit
# mixture-samples or resamples its posterior). Chain diagnostics (Rhat,
# autocorrelation-ESS) are only meaningful for the former: on i.i.d. draws
# split-Rhat sits at ~1 and ESS ~ n_draws by construction, which reads as a
# clean convergence result while saying nothing about approximation bias. The
# kind is carried explicitly on the fit (`$draws_kind`, stamped by
# `tulpa_dispatch`) or derived from the backend's registry `emits` property.

# Posterior representation kind: "chain", "iid", "point", or NA when unknown.
# A model-package wrapper (e.g. a tobs_fit) may leave its own `draws_kind`
# unset while carrying a nested-Laplace `joint_fit` that stamps the kind; that
# inner tag is read when the top-level tag and backend give no answer.
.tulpa_draws_kind <- function(fit) {
  k <- fit$draws_kind
  if (!is.null(k)) return(k)
  b <- fit$backend
  if (!is.null(b) && !is.null(BACKEND_REGISTRY[[b]])) {
    return(BACKEND_REGISTRY[[b]]$emits %||% NA_character_)
  }
  jk <- fit$joint_fit$draws_kind
  if (!is.null(jk)) return(jk)
  NA_character_
}

# TRUE when MCMC chain diagnostics describe the fit. Unknown provenance counts
# as a chain, so a fit that predates the `emits` tag is never silently refused.
.tulpa_is_chain <- function(fit) {
  k <- .tulpa_draws_kind(fit)
  is.na(k) || identical(k, "chain")
}

# Fit-time floor on a chain's mixing, below which the draws are not a usable
# posterior sample at all: split-Rhat above 1.05 or bulk-ESS under 100 on any
# main (non-indexed) parameter. Looser than check_diagnostics()'s 1.01 / 400 on
# purpose -- that is the bar for reporting, this is the bar for "the numbers
# printed are not yet the posterior" -- so a routine fit does not warn and an
# unconverged one cannot pass silently. A chain with bulk-ESS 2 printed clean
# intervals, and only an opt-in diagnostic said otherwise (gcol33/tulpa#875,
# #878).
.TULPA_FIT_RHAT_MAX <- 1.05
.TULPA_FIT_ESS_MIN  <- 100

# Check a fitted chain against that floor. Records the verdict on the fit
# (`$convergence`, which print() reads) and warns naming the offending
# parameters. A fit that is not a chain, or is too short to diagnose, passes
# through untouched.
.tulpa_check_fit_convergence <- function(fit, caller = "tulpa()") {
  if (!is.list(fit) || !.tulpa_is_chain(fit)) return(fit)
  diag <- tryCatch(suppressWarnings(suppressMessages(
    diagnostics(fit, measures = c("rhat", "ess_bulk")))),
    error = function(e) NULL)
  if (is.null(diag) || !NROW(diag)) return(fit)
  diag <- diag[diag$parameter %in% select_main_params(diag$parameter), ,
               drop = FALSE]
  bad_r <- is.finite(diag$rhat) & diag$rhat > .TULPA_FIT_RHAT_MAX
  bad_e <- is.finite(diag$ess_bulk) & diag$ess_bulk < .TULPA_FIT_ESS_MIN
  fit$convergence <- list(
    ok = !any(bad_r | bad_e),
    rhat_max = suppressWarnings(max(diag$rhat, na.rm = TRUE)),
    ess_bulk_min = suppressWarnings(min(diag$ess_bulk, na.rm = TRUE)),
    parameters = diag$parameter[bad_r | bad_e])
  if (fit$convergence$ok) return(fit)
  worst <- diag[bad_r | bad_e, , drop = FALSE]
  worst <- worst[order(worst$ess_bulk), , drop = FALSE]
  shown <- utils::head(worst, 5L)
  warning(sprintf(paste0(
    "%s: the '%s' chain has not converged -- %s. Its estimates and intervals ",
    "are not yet the posterior. Run longer (control$n_iter / control$warmup), ",
    "or use mode = 'hmc'; see check_diagnostics()."),
    caller, fit$backend %||% "sampler",
    paste(sprintf("%s: Rhat %.2f, bulk-ESS %.0f", shown$parameter, shown$rhat,
                  shown$ess_bulk), collapse = "; ")),
    call. = FALSE)
  fit
}

# Shared explanation for why chain diagnostics are withheld on a non-chain fit.
.tulpa_non_chain_msg <- function(fit) {
  k <- .tulpa_draws_kind(fit)
  b <- fit$backend %||% "this backend"
  if (identical(k, "point")) {
    return(sprintf("Backend '%s' returns a point summary (mode + covariance), not a posterior sample.", b))
  }
  sprintf(paste0(
    "Backend '%s' returns i.i.d. / resampled draws, not an MCMC chain: Rhat is ",
    "vacuous and ESS = n_draws by construction, so they do not diagnose ",
    "approximation accuracy. Assess this fit with the approximation/coverage ",
    "diagnostics instead (the IMH/Gibbs debias step, PIT residuals, recovery ",
    "checks)."), b)
}

#' Posterior parameter sample from a fit
#'
#' Returns a fit's posterior draws for summary purposes (quantiles, derived
#' quantities, density plots), regardless of how they were produced -- an MCMC
#' chain, a nested-Laplace node mixture, or a variational sample all answer
#' here. For the chain-only view used by convergence diagnostics, see
#' [mcmc_draws()].
#'
#' A fit that carries no draws returns `NULL` with a message naming its backend,
#' the posterior representation it carries instead, and the accessor that turns
#' that representation into draws where one exists -- an
#' absent draws matrix is a property of the backend, not a failure of the
#' accessor, and saying which is the difference between a diagnosable answer and
#' a bare `NULL`.
#'
#' @param fit A `tulpa_fit` (or subclass) carrying posterior `$draws`.
#' @return The posterior draws matrix/array, or `NULL` if the fit carries none.
#' @seealso [mcmc_draws()], [tulpa_posterior_draws()], [diagnostics()]
#' @examples
#' \donttest{
#' set.seed(1)
#' df <- data.frame(x = rnorm(80))
#' df$y <- rpois(80, exp(0.5 + 0.3 * df$x))
#' fit <- tulpa(y ~ x, data = df, family = "poisson")
#' dim(posterior_sample(fit))
#' }
#' @export
posterior_sample <- function(fit) {
  draws <- .fit_draws(fit)
  if (is.null(draws)) {
    message(.tulpa_no_draws_note(fit, "posterior_sample"))
    return(NULL)
  }
  draws
}

# What to say about a fit that carries no `$draws`. Names the backend, the
# posterior representation the fit DOES carry, and the accessor that samples it
# where there is one -- the same standard the Pareto-k / gamma_3 declines are
# held to: a withheld answer records why.
#' @keywords internal
.tulpa_no_draws_note <- function(fit, caller) {
  backend <- if (is.list(fit)) fit$backend %||% NA_character_ else NA_character_
  who <- if (is.na(backend)) {
    paste0("this ", paste(class(fit), collapse = "/"), " fit")
  } else paste0("the '", backend, "' fit")

  carries <- NULL
  if (is.list(fit)) {
    if (!is.null(fit$Q_csc_p_per_grid) && !is.null(fit$modes)) {
      carries <- paste(
        "the outer-grid Gaussian mixture over the full latent vector",
        "(`$modes` / `$Q_csc_*_per_grid` / `$weights`); draw from it with",
        "tulpa_posterior_draws(fit)")
    } else if (!is.null(fit$grid_hessians) && !is.null(fit$grid_modes)) {
      carries <- paste(
        "the outer-grid Gaussian mixture over the fixed effects",
        "(`$grid_modes` / `$grid_hessians` / `$weights`); draw from it with",
        "tulpa_posterior_draws(fit)")
    } else if (is.matrix(fit$modes) && !is.null(fit$weights)) {
      carries <- paste(
        "per-grid-cell modes and integration weights (`$modes` / `$weights`)",
        "but no per-cell covariance; refit with",
        "`control$keep_grid_hessians = TRUE` to make it samplable")
    } else if (!is.null(fit$mode) && !is.null(fit$H_beta)) {
      carries <- paste("a Laplace mode and its precision (`$mode` /",
                       "`$H_beta`), summarized by summary() / confint()")
    } else if (!is.null(fit$means) || is.matrix(fit$cov)) {
      carries <- paste("posterior moments (`$means` / `$cov`), summarized by",
                       "summary() / confint()")
    }
  }
  if (is.null(carries)) {
    carries <- "no posterior representation this accessor can read"
  }
  # A chain-stamped fit is a sampler's output, so an absent sample there is a
  # sample that was not kept rather than a representation without one.
  chain <- is.list(fit) && identical(fit[["draws_kind"]], "chain")
  sprintf("%s(): %s carries no posterior draws.%s It carries %s.",
          caller, who,
          if (chain) " It is stamped as an MCMC chain, which it did not retain." else "",
          carries)
}

# The one place a fit's posterior draws are read off the object.
#
# `[[` is deliberate, and load-bearing: `$` partial-matches on lists, so on a
# fit with no draws `fit$draws` silently resolves to `draws_kind` -- the string
# "iid" -- and every caller then treats a tag as a draws matrix. Gaussian-
# approximation fits (laplace, eb) carry `draws_kind` and no `draws`, which is
# exactly the case the draws accessors exist to answer NULL for.
#' @keywords internal
.fit_draws <- function(fit) {
  if (!is.list(fit)) return(NULL)
  draws <- fit[["draws"]] %||% fit[["samples"]]
  # A zero-row draws matrix holds no sample, so it is read as none.
  if (!is.null(draws) && NROW(draws) == 0L) return(NULL)
  draws
}

# Exact first two moments of a Gaussian mixture a draw matrix was sampled from.
#
# A sampler that draws i.i.d. from a known outer-grid mixture stamps it on its
# output as `attr(draws, "grid_mixture")`: `weights` (one per cell), `modes`
# (cells x p, the leading p columns of the draws) and the within-cell spread as
# either `vars` (cells x p marginal variances) or `covs` (a list of p x p
# covariances; a NULL or non-finite entry is a cell drawn at its mode). The
# moments follow from the law of total variance,
#   mean = sum_k w_k m_k,   var = sum_k w_k (v_k + m_k^2) - mean^2,
# so a summary can report them instead of the Monte-Carlo estimate the draws
# give. A draw matrix carries the attribute only while it IS that mixture's
# sample: draws the subspace debias corrected are a different posterior and are
# not stamped. NULL when there is no mixture.
#' @keywords internal
.grid_mixture_moments <- function(draws) {
  gm <- attr(draws, "grid_mixture")
  if (is.null(gm) || is.null(gm$modes) || is.null(gm$weights)) return(NULL)
  modes <- as.matrix(gm$modes)
  w <- as.numeric(gm$weights)
  if (length(w) != nrow(modes)) return(NULL)
  w[!is.finite(w) | w < 0] <- 0
  if (!any(w > 0)) return(NULL)
  w <- w / sum(w)
  p <- ncol(modes)
  vars <- if (!is.null(gm$vars)) {
    as.matrix(gm$vars)
  } else if (!is.null(gm$covs) && length(gm$covs) == nrow(modes)) {
    t(vapply(gm$covs, function(C) {
      d <- if (is.null(C)) rep(0, p) else diag(as.matrix(C))[seq_len(p)]
      if (all(is.finite(d))) d else rep(0, p)
    }, numeric(p)))
  } else return(NULL)
  if (!identical(dim(vars), dim(modes))) return(NULL)
  on <- w > 0
  wk <- w[on]
  mk <- modes[on, , drop = FALSE]
  m  <- colSums(wk * mk)
  v  <- colSums(wk * (vars[on, , drop = FALSE] + mk^2)) - m^2
  list(mean = m, sd = sqrt(pmax(v, 0)))
}

# Per-column posterior moments a fit carries EXACTLY for its draw matrix, as
# `mean` / `sd` vectors over the draw columns, NA where only the draws describe
# that column. The leading columns come from the stamped mixture
# (`.grid_mixture_moments()`); columns past it -- a model package's
# hyperparameter tail, drawn from its Gaussian conditional given the mixture
# block -- come from the fit's own `$means` / `$sds`, joined by parameter name.
# Both are read only on a stamped draw matrix, since that stamp is what says
# the draws were sampled from the moments the fit reports. NULL when nothing is
# exact.
#' @keywords internal
.tulpa_exact_moments <- function(draws, nm, fit = NULL) {
  mom <- .grid_mixture_moments(draws)
  if (is.null(mom)) return(NULL)
  n_col <- length(nm)
  mean <- rep(NA_real_, n_col)
  sd   <- rep(NA_real_, n_col)
  lead <- seq_len(min(length(mom$mean), n_col))
  mean[lead] <- mom$mean[lead]
  sd[lead]   <- mom$sd[lead]
  tail <- setdiff(seq_len(n_col), lead)
  fm <- if (is.list(fit)) fit[["means"]] else NULL
  fs <- if (is.list(fit)) fit[["sds"]] else NULL
  if (length(tail) && !is.null(names(fm)) && !is.null(names(fs))) {
    hit <- tail[nm[tail] %in% names(fm) & nm[tail] %in% names(fs)]
    mean[hit] <- as.numeric(fm[nm[hit]])
    sd[hit]   <- as.numeric(fs[nm[hit]])
  }
  ok <- is.finite(mean) & is.finite(sd)
  mean[!ok] <- NA_real_
  sd[!ok]   <- NA_real_
  if (!any(ok)) return(NULL)
  list(mean = mean, sd = sd)
}

#' MCMC chain draws from a fit
#'
#' Returns a fit's posterior draws only when they form a genuine MCMC chain
#' (`$draws_kind == "chain"`, or an untagged legacy fit); for an i.i.d. /
#' approximation fit (nested Laplace, VI, SMC, ...) it returns `NULL`, because
#' chain diagnostics do not apply. This is the accessor [diagnostics()]
#' gates on. For the provenance-agnostic posterior sample used by summaries,
#' see [posterior_sample()].
#'
#' @param fit A `tulpa_fit` (or subclass) carrying posterior `$draws`.
#' @return The chain draws matrix/array, or `NULL` for a non-chain fit.
#' @seealso [posterior_sample()], [diagnostics()]
#' @examples
#' \donttest{
#' set.seed(1)
#' d <- data.frame(x = rnorm(60))
#' d$y <- rpois(60, exp(0.5 + 0.3 * d$x))
#' fit <- tulpa(y ~ x, data = d, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 400L, warmup = 200L, n_chains = 2L,
#'                             seed = 1L))
#' dim(mcmc_draws(fit))
#' }
#' @export
mcmc_draws <- function(fit) {
  if (!.tulpa_is_chain(fit)) return(NULL)
  .fit_draws(fit)
}

# Pooled [total_draws x n_par] draws matrix with parameter colnames, from either
# a 2-D pooled matrix or a 3-D [iter, chain, param] array. The per-parameter
# diagnostic plots (ACF, Geweke) index draws[, param] and read colnames(), which
# a bare 3-D array breaks (NULL colnames, wrong slice); this normalizes both.
#' @keywords internal
.tulpa_pooled_draws <- function(fit) {
  draws <- .fit_draws(fit)
  if (is.null(draws)) return(NULL)
  if (length(dim(draws)) == 3L) {
    d  <- dim(draws)
    pn <- .tulpa_draw_names(dimnames(draws)[[3L]], fit, d[3L])
    # [iter, chain, param] is column-major with iter fastest, so reshaping to
    # nrow = iter * chain stacks each parameter's pooled draws into one column.
    m <- matrix(as.numeric(draws), nrow = d[1L] * d[2L], ncol = d[3L])
    colnames(m) <- pn
    return(m)
  }
  cn <- .tulpa_draw_names(colnames(draws), fit, ncol(draws))
  if (!identical(colnames(draws), cn)) colnames(draws) <- cn
  draws
}

# --- Draws extraction -------------------------------------------------------

# Split a fit's posterior draws into a list of per-chain [n_draws x n_par]
# matrices. Handles a 3D [iter, chain, param] array, a pooled matrix with a
# `chain_id` row map or an `n_chains` count (contiguous chain-major blocks),
# and the single-chain fallback.
.tulpa_chain_list <- function(fit) {
  draws <- .fit_draws(fit)
  if (is.null(draws)) return(NULL)

  if (length(dim(draws)) == 3L) {
    pn <- dimnames(draws)[[3L]]
    return(lapply(seq_len(dim(draws)[2L]), function(m) {
      mm <- draws[, m, , drop = TRUE]
      if (is.null(dim(mm))) mm <- matrix(mm, ncol = dim(draws)[3L])
      colnames(mm) <- pn
      mm
    }))
  }

  draws <- as.matrix(draws)
  cid <- .tulpa_chain_id(fit, nrow(draws))
  lapply(sort(unique(cid)), function(k) draws[cid == k, , drop = FALSE])
}

# Chain index (1..n_chains) of each of `n_rows` pooled draw rows: the fit's
# `chain_id` row map when it covers every row, else contiguous chain-major
# blocks of an `n_chains` count that divides the rows, else one chain.
.tulpa_chain_id <- function(fit, n_rows) {
  cid <- fit[["chain_id"]]
  nch <- fit[["n_chains"]] %||% 1L
  if (!is.null(cid) && length(cid) == n_rows) {
    return(as.integer(factor(cid)))
  }
  if (nch > 1L && n_rows %% nch == 0L) {
    return(rep(seq_len(nch), each = n_rows %/% nch))
  }
  rep(1L, n_rows)
}

#' Posterior draws as a 3D array
#'
#' Assembles a fitted model's posterior draws into an `[iteration, chain,
#' parameter]` array -- the layout expected by `bayesplot` and analogous to
#' `posterior::as_draws_array()`. Multiple chains are recognised from a 3D
#' draws array, a `$chain_id` row map, or an `$n_chains` count over chain-major
#' rows; a single pooled chain yields a one-chain array.
#'
#' As with [posterior_sample()], a fit that carries no draws returns `NULL` with
#' a message naming its backend and the representation it carries instead.
#'
#' @param fit A `tulpa_fit` (or subclass) carrying posterior `$draws`.
#' @return A numeric array with dimensions `[n_iter, n_chain, n_param]` and the
#'   parameter names on the third dimension, or `NULL` if the fit carries no
#'   draws. Chains of unequal length are truncated to the shortest.
#' @seealso [posterior_sample()], [tulpa_posterior_draws()], [diagnostics()]
#' @examples
#' \donttest{
#' set.seed(1)
#' d <- data.frame(x = rnorm(60))
#' d$y <- rpois(60, exp(0.5 + 0.3 * d$x))
#' fit <- tulpa(y ~ x, data = d, family = "poisson", mode = "hmc",
#'              control = list(n_iter = 400L, warmup = 200L, n_chains = 2L,
#'                             seed = 1L))
#' dim(tulpa_draws_array(fit))   # [iteration, chain, parameter]
#' }
#' @export
tulpa_draws_array <- function(fit) {
  arr <- .tulpa_draws_array(fit)
  if (is.null(arr)) message(.tulpa_no_draws_note(fit, "tulpa_draws_array"))
  arr
}

# Column names for a draws matrix, in the order they are trustworthy: the
# matrix's own, then the fit's `$param_names`, then positional. Several backends
# (mala, imh_laplace, pathfinder, vi) store an UNNAMED draws matrix alongside a
# fully populated `$param_names`, so reading colnames() alone reported every row
# as param1..paramN and a warning list of "Rhat > 1.01: param3" could not be
# traced back to a coefficient without counting columns by hand
# (gcol33/tulpa#714).
#' @keywords internal
.tulpa_draw_names <- function(nm, fit, p) {
  if (!is.null(nm) && length(nm) == p && all(nzchar(nm))) return(nm)
  base <- if (!is.null(nm) && length(nm) == p) nm else fit$param_names
  .tulpa_fill_names(base, p, "param")
}

# The array assembly without the user-facing note, for the internal callers that
# probe for draws and already have a designed fallback (the `posterior` interop
# conversion, k-fold, power-scaling). One body, two doors: the front door reports
# an absent representation, the probe does not narrate its own fallback.
#' @keywords internal
.tulpa_draws_array <- function(fit) {
  chain_list <- .tulpa_chain_list(fit)
  if (is.null(chain_list)) return(NULL)
  n_iter <- min(vapply(chain_list, nrow, integer(1)))
  p  <- ncol(chain_list[[1L]])
  nm <- .tulpa_draw_names(colnames(chain_list[[1L]]), fit, p)
  arr <- array(
    NA_real_, dim = c(n_iter, length(chain_list), p),
    dimnames = list(NULL, NULL, nm)
  )
  for (m in seq_along(chain_list)) {
    arr[, m, ] <- chain_list[[m]][seq_len(n_iter), , drop = FALSE]
  }
  arr
}

# bayesplot-shaped accessor: list(draws = <3D array>).
get_draws_array <- function(fit) list(draws = tulpa_draws_array(fit))

# Per-parameter chain diagnostics on a fit whose draws are a genuine MCMC
# chain. The provenance gate lives in `diagnostics()`; this builds the table.
#' @keywords internal
.tulpa_chain_diag_table <- function(fit, pars = NULL,
                                    measures = c("rhat", "ess_bulk", "ess_tail"),
                                    probs = c(0.05, 0.95)) {
  chain_list <- .tulpa_chain_list(fit)
  if (is.null(chain_list) || nrow(chain_list[[1L]]) < 4L) return(NULL)

  p  <- ncol(chain_list[[1L]])
  nm <- .tulpa_draw_names(colnames(chain_list[[1L]]), fit, p)
  keep <- if (is.null(pars)) seq_len(p) else which(nm %in% pars)
  if (length(keep) == 0L) return(NULL)

  n     <- nrow(chain_list[[1L]])
  cols  <- .tulpa_measure_cols(measures, probs)
  out   <- matrix(NA_real_, nrow = length(keep), ncol = length(cols),
                  dimnames = list(NULL, cols))
  for (i in seq_along(keep)) {
    j <- keep[i]
    sims <- vapply(chain_list, function(cd) cd[, j], numeric(n))
    if (is.null(dim(sims))) sims <- matrix(sims, ncol = length(chain_list))
    if (stats::sd(c(sims)) == 0) next                  # constant: leave NA
    vals <- tryCatch(.tulpa_measure_eval(sims, measures, probs),
                     error = function(e) NULL)
    if (!is.null(vals)) out[i, names(vals)] <- vals
  }
  data.frame(parameter = nm[keep], out, check.names = FALSE,
             stringsAsFactors = FALSE, row.names = NULL)
}

#' Select the "main" model parameters for diagnostic display
#'
#' Drops per-element latent-field entries (names ending in a bracketed index,
#' e.g. `u[12]`, `w[3]`) so plots and summaries focus on scalar coefficients
#' and hyperparameters rather than thousands of latent values.
#'
#' @param param_names Character vector of parameter names.
#' @return The subset of `param_names` that are not bracketed-index entries
#'   (or the full vector if that would be empty).
#' @examples
#' select_main_params(c("beta[1]", "beta[2]", "log_sigma_re", "re[1]", "re[2]"))
#' @export
select_main_params <- function(param_names) {
  keep <- !grepl("\\[[0-9]+(,[0-9]+)*\\]$", param_names)
  if (!any(keep)) return(param_names)
  param_names[keep]
}