R/nested_laplace_joint_helpers.R

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

# Grid / arm / copy / layout helpers.
#
# Part of the joint nested-Laplace driver; the public entry point
# tulpa_nested_laplace_joint() lives in nested_laplace_joint.R.

# --- helpers -----------------------------------------------------------------

# Cartesian product over a named list of spatial axes plus an optional
# alpha axis (copy coefficient) and optional per-arm phi axes.
# `phi_axes` is a list keyed by arm name; entries are either NULL/empty (no
# axis for that arm) or numeric vectors that become a new outer-grid axis
# named `phi_<arm>`. Returns a named list of paired vectors of identical
# length, ready to feed the C++ kernel. Phi axes vary slowest (added last)
# so within-spatial-block warm-starts stay good.
.joint_cartesian <- function(axes, has_copy, alpha_axis, phi_axes = NULL) {
    full <- if (has_copy) c(axes, list(alpha = alpha_axis)) else axes
    if (!is.null(phi_axes)) {
        active <- phi_axes[vapply(phi_axes, length, integer(1)) > 0L]
        if (length(active) > 0L) {
            names(active) <- paste0("phi_", names(active))
            full <- c(full, active)
        }
    }
    gr <- do.call(expand.grid,
                  c(full, list(KEEP.OUT.ATTRS = FALSE,
                                stringsAsFactors = FALSE)))
    out <- as.list(gr)
    if (!has_copy) out$alpha <- numeric(0)
    out
}

# Append any `phi_<arm>` columns from `grids` onto a backend's spatial
# theta_grid matrix so downstream posterior-moment helpers see phi as a
# regular hyperparameter axis.
.append_phi_columns <- function(base, grids) {
    phi_cols <- grep("^phi_", names(grids), value = TRUE)
    if (length(phi_cols) == 0L) return(base)
    extra <- do.call(cbind, lapply(phi_cols, function(c) {
        out <- as.numeric(grids[[c]]); attr(out, "name") <- c; out
    }))
    colnames(extra) <- phi_cols
    cbind(base, extra)
}

# Build the `phi_grid_per_arm` argument for the C++ kernels from a
# Cartesian-product `grids` list and arm names. Returns a list of length
# `n_arms`: entry k is either `NULL` (no phi axis for that arm -- kernel
# uses the parse-time scalar phi) or a NumericVector of length n_grid
# matching the flat outer-grid size. Phi columns in `grids` follow the
# `phi_<arm_name>` convention produced by `.joint_cartesian`.
.joint_phi_grid_per_arm <- function(grids, arm_names) {
    out <- vector("list", length(arm_names))
    any_active <- FALSE
    for (k in seq_along(arm_names)) {
        col <- paste0("phi_", arm_names[k])
        if (!is.null(grids[[col]])) {
            out[[k]] <- as.numeric(grids[[col]])
            any_active <- TRUE
        }
    }
    if (!any_active) NULL else out
}

# Compute a tile partition for the outer-grid loop's three-tier warm-start.
# A tile groups all
# outer-grid cells that share every hyperparameter coordinate except the
# copy coefficient alpha -- for the joint copy block under the (sigma, alpha)
# reparam, the shared latent prior Q and the donor-arm linear predictor are
# tile-constant, so the joint mode varies smoothly across the alpha axis
# within a tile. Using the tile's median-alpha cell as a warm-start for the
# remaining alpha cells saves 1-2 Newton iters each.
#
# Arguments:
#   non_alpha_matrix - numeric matrix [n_grid x n_axes_excl_alpha]; row k
#                      is the non-alpha coordinates of cell k.
#   alpha_vec        - numeric [n_grid]; alpha value for each cell.
#   n_grid           - integer; number of outer-grid cells (= nrow).
#
# Returns NULL when the partition has no useful structure (every cell its
# own tile, or only one tile total); otherwise a list with 0-based
# integer fields:
#   tile_ids[k]        - tile membership of cell k.
#   tile_pilot_cells[t]- cell index used as tile t's representative pilot.
#                        When the global pilot cell (n_grid %/% 2) falls
#                        into tile t, it is reused so the Tier-2 pass
#                        does not re-solve that tile.
.joint_compute_tile_partition <- function(non_alpha_matrix, alpha_vec, n_grid) {
    if (is.null(non_alpha_matrix) || ncol(non_alpha_matrix) == 0L) return(NULL)
    if (nrow(non_alpha_matrix) != n_grid || length(alpha_vec) != n_grid) {
        return(NULL)
    }
    keys <- do.call(paste, c(
        lapply(seq_len(ncol(non_alpha_matrix)), function(j) {
            formatC(non_alpha_matrix[, j], digits = 15L, format = "g")
        }),
        list(sep = "\r")
    ))
    uniq_keys <- unique(keys)
    tile_ids <- match(keys, uniq_keys) - 1L  # 0-based
    n_tiles <- length(uniq_keys)
    if (n_tiles <= 1L || n_tiles >= n_grid) return(NULL)

    # Tier-1 global pilot cell matches the C++ driver's `n_grid / 2`.
    k_global_pilot <- as.integer(n_grid %/% 2L)
    tile_of_global <- tile_ids[k_global_pilot + 1L]

    tile_pilot_cells <- integer(n_tiles)
    for (t in seq_len(n_tiles) - 1L) {
        if (t == tile_of_global) {
            # Global pilot doubles as this tile's pilot -- no Tier-2 solve.
            tile_pilot_cells[t + 1L] <- k_global_pilot
            next
        }
        cells_1based <- which(tile_ids == t)
        med <- stats::median(alpha_vec[cells_1based])
        best <- cells_1based[which.min(abs(alpha_vec[cells_1based] - med))][1L]
        tile_pilot_cells[t + 1L] <- as.integer(best - 1L)  # 0-based
    }
    list(tile_ids        = as.integer(tile_ids),
         tile_pilot_cells = as.integer(tile_pilot_cells))
}

# Normalise the user-facing `phi_grid` argument into a list keyed by arm
# name, with NULL entries for arms without a phi axis. Accepts either a
# named list (subset of arm names) or a positional list of length n_arms.
# Single-element entries are treated as no-axis (the parse-time scalar phi
# already serves as that arm's dispersion).
.normalise_phi_grid <- function(phi_grid, arm_names) {
    if (is.null(phi_grid)) return(NULL)
    if (!is.list(phi_grid)) {
        stop("`phi_grid` must be a list (named by arm or positional).",
             call. = FALSE)
    }
    out <- vector("list", length(arm_names))
    names(out) <- arm_names
    if (!is.null(names(phi_grid))) {
        unknown <- setdiff(names(phi_grid), arm_names)
        if (length(unknown) > 0L) {
            stop("`phi_grid` names not in `responses`: ",
                 paste(shQuote(unknown), collapse = ", "), ".", call. = FALSE)
        }
        for (nm in names(phi_grid)) {
            v <- phi_grid[[nm]]
            if (!is.null(v) && length(v) > 1L) out[[nm]] <- as.numeric(v)
        }
    } else {
        if (length(phi_grid) != length(arm_names)) {
            stop("positional `phi_grid` must have length n_arms (",
                 length(arm_names), ").", call. = FALSE)
        }
        for (k in seq_along(phi_grid)) {
            v <- phi_grid[[k]]
            if (!is.null(v) && length(v) > 1L) out[[k]] <- as.numeric(v)
        }
    }
    out
}

# Validate one arm spec and fill in defaults. `spatial_idx` selects the
# policy for the arm-level spatial index: "required" (single-block path:
# a mandatory field) or "optional" (multi-block path: per-arm idx vectors
# live inside each block spec, so a missing spatial_idx gets a length-N
# placeholder of zeros -- it satisfies parse_joint_arms' length check
# without contributing to eta).
.normalise_joint_arm_core <- function(a, k,
                                      spatial_idx = c("required", "optional")) {
    spatial_idx <- match.arg(spatial_idx)
    if (!is.list(a)) {
        stop("Arm ", k, ": expected a list of arm spec fields.", call. = FALSE)
    }
    must_have <- if (spatial_idx == "required") c("y", "X", "spatial_idx", "family")
                 else c("y", "X", "family")
    missing <- setdiff(must_have, names(a))
    if (length(missing)) {
        stop("Arm ", k, ": missing fields ", paste(shQuote(missing), collapse = ", "),
             ".", call. = FALSE)
    }
    N <- length(a$y)
    a$y <- as.numeric(a$y)
    if (!is.matrix(a$X)) {
        stop("Arm ", k, ": `X` must be a numeric matrix.", call. = FALSE)
    }
    if (nrow(a$X) != N) {
        stop("Arm ", k, ": nrow(X) (", nrow(a$X), ") must equal length(y) (",
             N, ").", call. = FALSE)
    }
    if (spatial_idx == "optional" && is.null(a$spatial_idx)) {
        a$spatial_idx <- rep(0L, N)
    }
    if (length(a$spatial_idx) != N) {
        stop("Arm ", k, ": length(spatial_idx) (", length(a$spatial_idx),
             ") must equal length(y) (", N, ").", call. = FALSE)
    }
    a$spatial_idx <- as.integer(a$spatial_idx)
    a$n_trials <- if (is.null(a$n_trials)) rep(1L, N) else as.integer(a$n_trials)
    if (length(a$n_trials) != N) {
        stop("Arm ", k, ": length(n_trials) (", length(a$n_trials),
             ") must equal length(y) (", N, ").", call. = FALSE)
    }
    # Optional grouped beta sufficient statistics: when
    # present, n_trials is the per-row group count and (slog_y, slog_1my) the
    # within-group sum of log(y) / log(1-y); the built-in beta spec reads them.
    if (!is.null(a$slog_y)) {
        a$slog_y   <- as.numeric(a$slog_y)
        a$slog_1my <- as.numeric(a$slog_1my)
        if (length(a$slog_y) != N || length(a$slog_1my) != N) {
            stop("Arm ", k, ": length(slog_y)/length(slog_1my) must equal ",
                 "length(y) (", N, ").", call. = FALSE)
        }
    }
    # Optional interval-censored Gaussian bounds (family == "interval_gaussian"):
    # row i records the latent value fell in (lower[i], upper[i]] on the predictor
    # scale, with -Inf / +Inf the open outer classes. The built-in interval spec
    # reads (lower, upper) in place of the point response y.
    if (!is.null(a$lower) || !is.null(a$upper)) {
        if (is.null(a$lower) || is.null(a$upper)) {
            stop("Arm ", k, ": `lower` and `upper` must be supplied together.",
                 call. = FALSE)
        }
        a$lower <- as.numeric(a$lower)
        a$upper <- as.numeric(a$upper)
        if (length(a$lower) != N || length(a$upper) != N) {
            stop("Arm ", k, ": length(lower)/length(upper) must equal ",
                 "length(y) (", N, ").", call. = FALSE)
        }
        if (any(is.na(a$lower)) || any(is.na(a$upper)) || any(a$lower >= a$upper)) {
            stop("Arm ", k, ": each `lower` must be finite-or-(-Inf), non-NA, ",
                 "and strictly below its `upper`.", call. = FALSE)
        }
    }
    # Optional upper-truncated Gaussian ceiling (family == "truncated_gaussian"):
    # row i's latent log-response is truncated to <= trunc_upper[i] on the predictor
    # scale; +Inf => no truncation on that row. The built-in truncated spec reads
    # (y, trunc_upper).
    if (!is.null(a$trunc_upper)) {
        a$trunc_upper <- as.numeric(a$trunc_upper)
        if (length(a$trunc_upper) != N) {
            stop("Arm ", k, ": length(trunc_upper) (", length(a$trunc_upper),
                 ") must equal length(y) (", N, ").", call. = FALSE)
        }
        if (any(is.na(a$trunc_upper))) {
            stop("Arm ", k, ": `trunc_upper` must be non-NA (use +Inf for an ",
                 "untruncated row).", call. = FALSE)
        }
    }
    a$re_idx <- if (is.null(a$re_idx)) rep(0, N) else as.numeric(a$re_idx)
    if (length(a$re_idx) != N) {
        stop("Arm ", k, ": length(re_idx) (", length(a$re_idx),
             ") must equal length(y) (", N, ").", call. = FALSE)
    }
    a$n_re_groups <- as.integer(a$n_re_groups %||% 0L)
    a$sigma_re    <- as.numeric(a$sigma_re %||% 1.0)
    a$family      <- as.character(a$family)
    a$phi         <- as.numeric(a$phi %||% 1.0)
    a <- .normalise_arm_field_coef(a, k)
    a <- .normalise_arm_cell_coupling(a, k, N)
    a
}

# Validate one arm spec and fill in defaults (single-block path:
# arm-level `spatial_idx` is required).
.normalise_joint_arm <- function(a, k) {
    .normalise_joint_arm_core(a, k, spatial_idx = "required")
}

# Parse per-arm `coupled` / `cell_obs_map`. When `coupled` is TRUE, the
# inner Newton routes this arm's per-cell contribution through the
# registered CellCouplingSpec; `cell_obs_map[i]` is the 1-based cell id
# for row i of this arm.
.normalise_arm_cell_coupling <- function(a, k, N) {
    coupled <- isTRUE(a$coupled)
    a$coupled <- coupled
    if (!coupled) {
        a$cell_obs_map <- integer(0)
        return(a)
    }
    m <- a$cell_obs_map
    if (is.null(m)) {
        stop("Arm ", k, ": coupled = TRUE requires `cell_obs_map`.",
             call. = FALSE)
    }
    m <- as.integer(m)
    if (length(m) != N) {
        stop("Arm ", k, ": length(cell_obs_map) (", length(m),
             ") must equal length(y) (", N, ").", call. = FALSE)
    }
    if (any(is.na(m)) || any(m < 1L)) {
        stop("Arm ", k, ": `cell_obs_map` entries must be positive integers.",
             call. = FALSE)
    }
    a$cell_obs_map <- m
    a
}

# Parse the per-arm `field_coef` spec into resolved fields. Each arm carries
# one of: numeric scalar (constant), character (named outer-grid hyperparam
# axis), or list(name = , grid = ) (embedded axis declaration). The default
# is `field_coef = 1` (donor behaviour -- existing non-copy arms).
#
# Resolves to two fields on the arm spec:
#   * `field_coef_const` numeric scalar (default 1) -- the constant per-arm
#     multiplier always applied to that arm's field amplitude. For arms with
#     a hyperparam-driven coefficient this is 1 (the hyperparam carries the
#     coefficient).
#   * `field_coef_axis` NULL or list(name = , grid = , n = ) -- when set, this
#     arm declares a hyperparam-driven field coefficient. The driver maps that
#     declaration to the existing `copy = list(arm, alpha_grid)` plumbing
#     (at most one such axis is supported). `grid` states the axis's nodes; `n`
#     re-reads the engine's own axis at a higher resolution instead
#     (gcol33/tulpa#633), keeping its atom at 0 and its slab bounds.
#     `grid_auto` records whether a stated `grid` carried the `auto_grid()`
#     marker, which is the provenance the refinement rule reads.
.normalise_arm_field_coef <- function(a, k) {
    fc <- a$field_coef
    if (is.null(fc)) {
        a$field_coef_const <- 1.0
        a$field_coef_axis  <- NULL
        return(a)
    }
    if (is.numeric(fc) && length(fc) == 1L) {
        if (!is.finite(fc) || fc < 0) {
            stop("Arm ", k, ": `field_coef` numeric scalar must be a finite ",
                 "non-negative number (got ", fc, ").", call. = FALSE)
        }
        a$field_coef_const <- as.numeric(fc)
        a$field_coef_axis  <- NULL
        return(a)
    }
    if (is.character(fc) && length(fc) == 1L) {
        a$field_coef_const <- 1.0
        a$field_coef_axis  <- list(name = fc, grid = NULL, alpha_n = NULL,
                                   grid_auto = FALSE)
        return(a)
    }
    if (is.list(fc)) {
        nm <- fc$name
        if (is.null(nm) || !is.character(nm) || length(nm) != 1L) {
            stop("Arm ", k, ": `field_coef` list must carry a single character ",
                 "`name`.", call. = FALSE)
        }
        gr <- fc$grid
        # `as.numeric()` drops attributes, so the `auto_grid()` marker -- a
        # wrapper package declaring these nodes a default of its own rather than
        # a user's choice -- is read off the value before it is coerced and
        # carried on the resolved axis, where the refinement provenance rule
        # reads it. A declaration asking for the nodes AS WRITTEN
        # (`auto_grid(place = FALSE)`) is a statement about where the fit
        # integrates like a pin is, so refinement treats it as one.
        gr_auto <- is_auto_grid(gr) && auto_grid_place(gr)
        if (!is.null(gr)) {
            gr <- as.numeric(gr)
            if (length(gr) == 0L || any(!is.finite(gr)) || any(gr < 0)) {
                stop("Arm ", k, ": `field_coef$grid` must be a non-empty ",
                     "numeric vector of finite non-negative values.",
                     call. = FALSE)
            }
        }
        # `[[` and not `$`: `$` PARTIAL-matches on a list, so `fc$n` resolves
        # to `fc$name` on every spec that names its coefficient and feeds a
        # character into the integer check below.
        an <- fc[["n"]]
        if (!is.null(an)) {
            an <- suppressWarnings(as.integer(an))
            if (length(an) != 1L || is.na(an) || an < 1L) {
                stop("Arm ", k, ": `field_coef$n` must be a single integer ",
                     ">= 1.", call. = FALSE)
            }
        }
        a$field_coef_const <- 1.0
        a$field_coef_axis  <- list(name = as.character(nm), grid = gr,
                                   alpha_n = an, grid_auto = gr_auto)
        return(a)
    }
    stop("Arm ", k, ": `field_coef` must be NULL, a numeric scalar, a single ",
         "character name, or `list(name = , grid = )`.", call. = FALSE)
}


# Decide which arm (if any) is the copy arm and what the alpha grid is.
# Single-block path: the copy coefficient's outer-grid axis lives on the arm
# that declares it.
#
# Source of truth: `arms` already carries the resolved per-arm `field_coef_axis`
# / `field_coef_const` from `.normalise_joint_arm`. When any arm declared a
# hyperparam axis, the resolved axis arm + grid are returned as the kernel-facing
# copy spec.
#
# `field_coef_const` carries the per-arm constant multipliers (default 1 for
# every arm). Arms with `field_coef = 0` get const 0; arms with `field_coef =
# c` get c. These multipliers ride alongside any hyperparam axis on the same
# arm (axis multiplied by const), so they are returned to the caller and
# threaded down to the kernel via `arms[k].field_coef` (see the C++ side).
.resolve_copy <- function(responses, prior, type) {
    n_arms <- length(responses)
    consts <- vapply(responses, function(a) {
        if (is.null(a$field_coef_const)) 1.0 else as.numeric(a$field_coef_const)
    }, numeric(1))
    axes <- lapply(responses, function(a) a$field_coef_axis)
    has_axis_per_arm <- !vapply(axes, is.null, logical(1))
    if (sum(has_axis_per_arm) > 1L) {
        stop("At most one arm may declare a hyperparam `field_coef` axis. Got ",
             sum(has_axis_per_arm), ".", call. = FALSE)
    }
    if (!any(has_axis_per_arm)) {
        return(list(has_copy = FALSE, copy_arm_zero = -1L,
                    alpha_grid = numeric(0),
                    field_coef_const = consts))
    }
    k_axis <- which(has_axis_per_arm)
    spec   <- axes[[k_axis]]
    alpha_axis <- .nl_copy_alpha_axis(spec[["grid"]], spec[["alpha_n"]],
                                      what = paste0("Arm ", k_axis))
    if (length(alpha_axis) == 0L) {
        stop("Arm ", k_axis, ": `field_coef$grid` must have at least one ",
             "non-negative value.", call. = FALSE)
    }
    if (any(alpha_axis < 0)) {
        stop("Arm ", k_axis, ": `field_coef$grid` values must be non-negative.",
             call. = FALSE)
    }
    list(has_copy = TRUE, copy_arm_zero = as.integer(k_axis - 1L),
         alpha_grid = as.numeric(alpha_axis),
         field_coef_const = consts)
}

# Weighted-quantile median + 2.5/97.5 empirical CI for every hyperparameter
# axis are produced generically by `.nl_axis_quantiles` (defined in
# nested_laplace.R) and attached by `.nl_posterior_moments`. No alpha-
# specific helper is needed: median/CI are the recommended summary for
# right-skewed scale-like axes (alpha = sigma_pos/sigma_occ, sigma,
# range, phi at small n), while mean +/- SD remains available on the
# same axis names via `theta_mean / theta_sd`.

# Alpha refinement piggybacks on the generic consistency pass: alpha
# appears in `.hyper_refinable_names`, and `.hyper_consistency_pass`
# fires whenever the alpha marginal spreads over too few nodes to carry a
# spread. No bespoke helper.


# --- which axes the joint grid carries, and how far each may be refined -------
#
# The axes one `grids` list produces specs for. The alpha column is present in
# the paired-vector representation whether or not a copy arm was declared, so a
# fit with no copy drops it.
.joint_spec_axis_names <- function(grids, cp) {
    axes <- names(grids)
    has_alpha <- isTRUE(cp$has_copy) && length(grids$alpha) > 0L
    if (!has_alpha) axes <- setdiff(axes, "alpha")
    axes
}

# Which axes the joint refinement machinery can place nodes on at all. This is
# a property of the DRIVER -- the field SD, the copy coefficient and the per-arm
# dispersions are the axes its passes know how to propose on -- not of where the
# nodes came from, which is the separate question `.joint_axis_is_stated()`
# answers.
#
# A multi-block grid prefixes the axes a block owns (`b2.sigma`) and leaves the
# per-arm dispersions bare (`phi_<arm>`), so a `phi_`-prefixed name is a
# dispersion only when no block prefix is present: `b1.phi_gp` is a block's own
# lengthscale. The field SD is `sigma` on a field-SD parameterisation and `tau`
# on a precision one, and the two are one axis under `tau = 1 / sigma^2`.
.JOINT_REFINABLE_BLOCK_AXES <- c("sigma", "tau", "alpha")

# Is `axis` a per-arm dispersion (`phi_<arm>`)? Only an unprefixed name is:
# `b1.phi_gp` is a block's own NNGP lengthscale.
.joint_axis_is_dispersion <- function(axis) {
    !grepl("^b[0-9]+[.]", axis) & startsWith(axis, "phi_")
}

#
# `dispersions = FALSE` is the registry door's single-block grid, whose axes are
# bare but name no per-arm dispersion (`phi_gp` is the nngp lengthscale there).
.joint_axis_refine_eligible <- function(axis, dispersions = TRUE) {
    if (grepl("^b[0-9]+[.]", axis)) {
        return(.hyper_axis_bare(axis) %in% .JOINT_REFINABLE_BLOCK_AXES)
    }
    axis %in% .JOINT_REFINABLE_BLOCK_AXES ||
        (isTRUE(dispersions) && .joint_axis_is_dispersion(axis))
}

# One entry of the user-facing `phi_grid` argument, by arm name. The argument is
# either named by arm or positional over `arm_names`, and the axis it produces
# is `phi_<arm name>` either way.
.joint_phi_grid_entry <- function(phi_grid, arm, arm_names = NULL) {
    if (!is.list(phi_grid) || !length(phi_grid)) return(NULL)
    if (!is.null(names(phi_grid))) {
        if (!arm %in% names(phi_grid)) return(NULL)
        return(phi_grid[[arm]])
    }
    k <- match(arm, arm_names %||% character(0))
    if (is.na(k) || k > length(phi_grid)) return(NULL)
    phi_grid[[k]]
}

# Did the CALLER write this axis's nodes down?
#
# A stated axis is a statement about where the fit integrates, so refinement
# resolves it more finely and does not leave it; an axis the engine placed
# carries no such statement. Three things count as engine-placed, matching the
# provenance vocabulary the recentring pass already uses
# (`.nl_axis_hold()`): no nodes given at all (the copy axis resolved from
# `alpha_n`, or from the engine default), nodes marked with `auto_grid()` by a
# wrapper package that computed a default of its own, and nodes that ARE the
# engine's own default axis, which carry nothing a statement would add. A
# package default marked `auto_grid(place = FALSE)` is NOT among them: asking
# for the nodes as written is the statement, whoever made it.
.joint_axis_is_stated <- function(axis, arms, phi_grid = NULL,
                                  arm_names = NULL, phi_auto = NULL,
                                  prior = NULL, grid_auto = logical(0)) {
    if (identical(axis, "sigma")) {
        return(!is.null(.nl_axis_hold(prior, "sigma_grid",
                                      .nl_auto_fields_at(grid_auto),
                                      type = ".joint_areal")))
    }
    if (identical(axis, "alpha")) {
        for (a in arms) {
            fc <- a$field_coef_axis
            if (is.null(fc)) next
            g <- fc$grid
            if (is.null(g) || !length(g)) return(FALSE)
            if (isTRUE(fc$grid_auto)) return(FALSE)
            return(!.nl_axis_matches_default(g, "alpha_grid", ".copy"))
        }
        return(FALSE)
    }
    .joint_phi_axis_is_stated(axis, phi_grid, arm_names, phi_auto)
}

# Did the caller write down the nodes of a per-arm dispersion axis
# (`phi_<arm>`)? Anything else is not this helper's business and reads FALSE.
.joint_phi_axis_is_stated <- function(axis, phi_grid = NULL, arm_names = NULL,
                                      phi_auto = NULL) {
    if (!startsWith(axis, "phi_")) return(FALSE)
    arm <- sub("^phi_", "", axis)
    v <- .joint_phi_grid_entry(phi_grid, arm, arm_names)
    if (is.null(v) || length(v) < 2L) return(FALSE)
    # The front door's record first: it is taken before the markers are
    # stripped, so it is the only reading available once they are gone. A
    # caller reaching this with the marked value itself is read off the
    # value.
    if (!is.null(phi_auto)) return(!is.null(.nl_phi_axis_hold(arm, phi_auto)))
    !is_auto_grid(v) || !auto_grid_place(v)
}

# The refinement mode of every axis of one joint grid, as a named character
# vector over `.NL_AXIS_REFINE_MODES`. Eligible axes take the provenance
# default, `is_stated(axis)` saying whether the caller wrote the axis's nodes
# down; everything else takes "none". `user` -- the caller's
# `control$axis_refine`, already validated -- overrides per axis. The one rule
# both joint drivers resolve their modes through, each supplying only how it
# reads provenance off its own call.
.joint_axis_refine_modes_by <- function(axes, is_stated, user = NULL,
                                        eligible = .joint_axis_refine_eligible) {
    out <- stats::setNames(rep("none", length(axes)), axes)
    for (a in axes) {
        if (!eligible(a)) next
        out[[a]] <- if (is_stated(a)) .nl_axis_refine("stated")
                    else .nl_axis_refine("placed")
    }
    if (!is.null(user)) {
        for (a in intersect(names(user), axes)) out[[a]] <- user[[a]]
    }
    out
}

.joint_axis_refine_modes <- function(grids, cp, arms, phi_grid = NULL,
                                     arm_names = NULL, user = NULL,
                                     phi_auto = NULL, prior = NULL,
                                     grid_auto = logical(0)) {
    .joint_axis_refine_modes_by(
        .joint_spec_axis_names(grids, cp),
        function(a) .joint_axis_is_stated(a, arms, phi_grid, arm_names,
                                          phi_auto, prior, grid_auto),
        user)
}

# The multi-block counterpart of `.joint_axis_is_stated()`: the same question,
# asked of the prior's blocks and the copy specs instead of one prior and the
# arms' `field_coef`. A block's field SD is stated when the block carries the
# grid field it is laid on -- `sigma_grid` or `tau_grid` -- as something other
# than the engine's own default (`.nl_axis_hold()`), compared against the axis
# the block's path defaults: a copy block's `sigma_grid` is the copy convention,
# any other block's the registry's. The copy scale reads the copy spec's
# `alpha_grid`, and a dispersion the front door's `phi_auto` record.
.joint_multi_axis_is_stated <- function(axis, prior, grid_auto, copy = NULL,
                                        copy_blocks = integer(0),
                                        phi_grid = NULL, arm_names = NULL,
                                        phi_auto = NULL) {
    if (!grepl("^b[0-9]+[.]", axis)) {
        return(.joint_phi_axis_is_stated(axis, phi_grid, arm_names, phi_auto))
    }
    b <- as.integer(sub("^b([0-9]+)[.].*$", "\\1", axis))
    bare <- .hyper_axis_bare(axis)
    if (identical(bare, "alpha")) {
        i <- .joint_copy_spec_for_block(copy, b)
        declared <- if (is.na(i)) NULL
                    else .joint_copy_specs(copy)[[i]]$alpha_grid
        return(!is.null(.nl_axis_hold(list(alpha_grid = declared), "alpha_grid",
                                      type = ".copy")))
    }
    !is.null(.nl_axis_hold(
        prior[[b]], paste0(bare, "_grid"), .nl_auto_fields_at(grid_auto, b),
        type = if (b %in% copy_blocks) ".copy"
               else tolower(prior[[b]]$type %||% "")))
}

# Validate `control$axis_refine` against the axes this fit actually has.
#
# Accepts a named character vector (or list) keyed by axis name, or one unnamed
# value applying to every eligible axis. An unknown axis name is an error rather
# than a silent no-op, and so is asking for nodes on an axis the driver's passes
# cannot place any on: a knob that quietly does nothing is how a caller comes to
# believe a range was honoured when it was not.
.joint_check_axis_refine <- function(x, axis_names,
                                     eligible_fn = .joint_axis_refine_eligible) {
    if (is.null(x)) return(NULL)
    if (is.list(x)) x <- unlist(x, use.names = TRUE)
    if (!is.character(x) || !length(x) || anyNA(x)) {
        stop("`control$axis_refine` must be a character vector of ",
             paste(sprintf('"%s"', .NL_AXIS_REFINE_MODES), collapse = " / "),
             ", named by outer-grid axis.", call. = FALSE)
    }
    bad <- setdiff(unique(unname(x)), .NL_AXIS_REFINE_MODES)
    if (length(bad)) {
        stop("`control$axis_refine`: unknown mode ",
             paste(shQuote(bad), collapse = ", "), ". Use one of ",
             paste(sprintf('"%s"', .NL_AXIS_REFINE_MODES), collapse = ", "),
             ".", call. = FALSE)
    }
    eligible <- axis_names[vapply(axis_names, eligible_fn, logical(1))]
    nm <- names(x)
    if (is.null(nm) || !all(nzchar(nm))) {
        if (length(x) != 1L) {
            stop("`control$axis_refine` must be named by axis, or a single ",
                 "mode applying to every refinable axis.", call. = FALSE)
        }
        if (!length(eligible)) return(NULL)
        return(stats::setNames(rep(unname(x), length(eligible)), eligible))
    }
    unknown <- setdiff(nm, axis_names)
    if (length(unknown)) {
        stop("`control$axis_refine` names axes this fit does not have: ",
             paste(shQuote(unknown), collapse = ", "), ". Its axes are ",
             paste(shQuote(axis_names), collapse = ", "), ".", call. = FALSE)
    }
    wrong <- nm[!nm %in% eligible & x != "none"]
    if (length(wrong)) {
        stop("`control$axis_refine`: ",
             paste(shQuote(wrong), collapse = ", "),
             " cannot take refinement nodes on this driver, so only \"none\" ",
             "applies there. The axes it refines are ",
             if (length(eligible)) paste(shQuote(eligible), collapse = ", ")
             else "none of this fit's", ".", call. = FALSE)
    }
    x
}

# The mode one axis ends up at. Absent provenance -- a spec list rebuilt from an
# assembled grid, which no refinement pass reads -- keeps the driver's
# eligibility as the answer.
.joint_axis_refine_mode <- function(axis_refine, axis) {
    if (!is.null(axis_refine) && axis %in% names(axis_refine)) {
        return(as.character(axis_refine[[axis]]))
    }
    if (.joint_axis_refine_eligible(axis)) .nl_axis_refine("placed") else "none"
}

# Generic axis-spec adapter (Step 3).
#
# Builds the list of `hyper_axis_spec` objects the generic refinement / consistency
# helpers (`R/hyper_grid_refine.R`) consume from the joint driver's paired-vector
# `grids` + copy state. The spec's `grid` field is informational (refinement reads
# the per-cell axis values from `theta_grid[, axis]` directly), but populating it
# with the sorted unique levels keeps the spec object self-describing.
#
# Hardcoded metadata captures what the legacy joint helpers `.axis_is_log_scale`,
# `.axis_bounds`, `.axis_refinement_order` encoded by name -- the same set, in
# one place. `axis_refine` carries the one piece of metadata that is NOT a
# property of the axis's name: how far each axis may be refined, which follows
# from where its nodes came from and so is resolved by
# `.joint_axis_refine_modes()` at the front door.
# `user_priors` carries the caller's regularizing hyperpriors, named by the axis
# each applies to (`sigma`, `alpha`, `phi` -- the last matching every `phi_<arm>`
# axis), matching how `.joint_hp_vec_for_grids()` folds them into
# `log_marginal`. An axis with one keeps the flat-over-declared-span measure,
# because applying the engine's declared density as well would put two priors on
# one axis, and carries the density on the spec so the axis quadrature knows
# what the fold will apply -- which is what lets a declared point mass keep its
# prior probability (`.hyper_atom_fold_scale()`).
.joint_axis_specs <- function(grids, cp, user_priors = NULL,
                              copy_atom_mass = .TULPA_COPY_ATOM_MASS,
                              copy_slab = "exponential",
                              axis_refine = NULL,
                              folded_axes = NULL) {
    copy_slab <- .hyper_check_copy_slab(copy_slab)
    axes <- .joint_spec_axis_names(grids, cp)
    lapply(axes, function(a) {
        # A multi-block grid prefixes its axes (`b1.sigma`), so the scale, the
        # domain and the point-mass metadata are resolved on the bare axis name:
        # an axis's support does not change because the block it is laid on is
        # the second one. Which axes refine, and in what order, stays keyed on
        # the full name.
        bare <- sub("^b[0-9]+[.]", "", a)
        scale_known <- .hyper_axis_scale(bare)
        log_scale <- isTRUE(scale_known)
        bounds <- .hyper_spec_bounds(bare)
        # How far refinement may move this axis: whether it participates at all,
        # and whether it may place a node past the nodes it was given. Both come
        # from `axis_refine`, resolved from the caller's provenance one level up
        # -- deciding by axis NAME would opt a caller's stated nodes into being
        # extended (gcol33/tulpa#658).
        mode <- .joint_axis_refine_mode(axis_refine, a)
        refine_priority <- if (identical(bare, "alpha")) 1L
                           else if (startsWith(a, "phi_")) 2L
                           else 100L
        # The copy scale carries an explicit zero level ("no coupling"), which
        # is a point mass rather than part of the log continuum, so it needs a
        # declared prior probability. Fixing it here keeps it independent of how
        # many continuum nodes the grid ends up with.
        spec <- hyper_axis_spec(
            name      = a,
            grid      = sort(unique(as.numeric(grids[[a]]))),
            log_scale = log_scale,
            bounds    = bounds,
            refinable = !identical(mode, "none"),
            extend    = identical(mode, "extend"),
            atom_mass = if (identical(bare, "alpha")) copy_atom_mass
        )
        spec$refine_priority <- refine_priority
        # An axis this table does not classify keeps equal node weights, which
        # is what the engine integrated before any coordinate was declared.
        if (is.na(scale_known)) spec$unweighted <- TRUE
        # A flat measure on a log axis is improper, so the support is a prior
        # choice. The incoming grid is what the user declared, so the support is
        # its span widened by half a node step at each end: that is the region
        # the initial nodes already tile, so an unrefined grid integrates what it
        # always did, and fixing it here, ahead of any refinement, is what keeps
        # the measure independent of the data.
        pos <- sort(unique(spec$grid[spec$grid > 0]))
        user_fn <- if (identical(a, "sigma")) user_priors$sigma
                   else if (identical(a, "alpha")) user_priors$alpha
                   else if (startsWith(a, "phi_")) user_priors$phi
                   else NULL
        user_prior <- !is.null(user_fn)
        if (user_prior) {
            # The joint driver's hyperprior families are densities on the axis's
            # natural scale.
            spec$log_prior <- user_fn
            spec$log_prior_coord <- "natural"
        }
        if (log_scale && length(pos) >= 2L) {
            copy_exp_slab <- identical(bare, "alpha") && !user_prior &&
                             identical(copy_slab, "exponential")
            if (copy_exp_slab) {
                # The copy scale is the axis a user cannot be expected to bracket
                # in advance, and the one carrying the point mass the continuum
                # is weighed against. It therefore gets a proper density over the
                # whole positive line rather than a span: refinement can then
                # follow the posterior out as far as it needs to, and still be
                # integrating the measure declared here, before the fit.
                #
                # `copy_slab = "flat"` takes the other branch instead, so the
                # copy scale is flat in log alpha over the span its declared
                # nodes tile, the same measure sigma and phi carry.
                spec$slab_log_density <- .hyper_copy_slab_density(max(pos))
            } else if (!a %in% folded_axes) {
                # An axis whose normalised density is folded into the marginal
                # (`R/hyperprior_default.R`) needs no span to make its measure
                # proper, and refinement may follow the posterior past its nodes.
                bd <- .hyper_default_coord_bounds(.hyper_axis_coord(pos, spec))
                spec$slab_bounds <- exp(bd)
            }
        }
        spec
    })
}

# Axis specs for a grid that is already assembled, keyed off its column names.
# The multi-block driver builds its grid before any spec list exists, so it
# recovers the same metadata from the columns rather than carrying a second
# description of the same axes. `logchol` is each free-covariance block's
# declared design (`.hp_logchol_designs()`), stamped on the block's specs as
# `logchol_design`; it defaults to the design `theta_grid` itself declares, and a
# caller measuring a grid it did not declare -- a subset of the declared tensor
# -- passes the declaration it folded the hyperprior under.
#' Rebuild axis specs from an already-assembled outer grid
#'
#' Rebuilds each axis's declared spec (support, log/linear coordinate,
#' prior) from a `theta_grid`'s columns, keyed off column names -- the same
#' metadata the multi-block joint driver recovers rather than carrying a
#' second description of the axes it already assembled. A column holding one
#' value across every cell is treated as a fixed setting rather than an axis
#' (e.g. a fixed NB-size node) and dropped. Intended for a caller that needs
#' an axis's declared span, coordinate or prior off a settled grid without
#' restating tulpa's own hyperprior declarations alongside it.
#'
#' @param theta_grid A named `[n_cells x n_axes]` matrix (see
#'   [tulpa_theta_matrix()]).
#' @param copy_slab `"exponential"` or `"flat"`; the copy-scale axis's
#'   continuum measure (see [tulpa_hyper_check_copy_slab()]).
#' @param folded_axes Optional names of axes folded onto `[0, Inf)`.
#' @param logchol Free-covariance block design(s) as returned by
#'   `.hp_logchol_designs()`; defaults to the design `theta_grid` itself
#'   declares.
#' @return A named list of per-axis spec lists, or `NULL` if `theta_grid` has
#'   no axis names.
#' @seealso [tulpa_theta_matrix()], [tulpa_grid_log_quad()],
#'   [tulpa_hyper_grid_supports()]
#' @examples
#' \donttest{
#' set.seed(1)
#' S <- 30L                                   # spatial units in a chain
#' nb <- lapply(seq_len(S), function(s) setdiff(c(s - 1L, s + 1L), c(0L, S + 1L)))
#' nn <- lengths(nb)
#' field <- as.numeric(scale(cumsum(rnorm(S, 0, 0.4))))
#' idx <- rep(seq_len(S), each = 6L); n <- length(idx); x <- rnorm(n)
#' y <- rbinom(n, 1L, plogis(-0.2 + 0.6 * x + field[idx]))
#' prior <- list(type = "icar", n_spatial_units = S, spatial_idx = idx,
#'               adj_row_ptr = c(0L, cumsum(nn)), adj_col_idx = unlist(nb) - 1L,
#'               n_neighbors = nn, tau_grid = c(0.5, 1, 2, 4, 8))
#' fit <- tulpa_nested_laplace(y, rep(1L, n), cbind(1, x), prior = prior,
#'                             family = "binomial",
#'                             control = list(progress = FALSE))
#' specs <- tulpa_joint_axis_specs_from_grid(tulpa_theta_matrix(fit))
#' specs[[1]]$log_scale
#' }
#' @export
tulpa_joint_axis_specs_from_grid <- function(
    theta_grid, copy_slab = "exponential", folded_axes = NULL,
    logchol = .hp_logchol_designs(theta_grid)) {
  .joint_axis_specs_from_grid(theta_grid, copy_slab = copy_slab,
                              folded_axes = folded_axes, logchol = logchol)
}

.joint_axis_specs_from_grid <- function(theta_grid,
                                        copy_slab = "exponential",
                                        folded_axes = NULL,
                                        logchol = .hp_logchol_designs(theta_grid),
                                        axis_refine = NULL) {
    if (is.null(theta_grid) || is.null(colnames(theta_grid))) return(NULL)
    theta_grid <- as.matrix(theta_grid)
    grids <- stats::setNames(
        lapply(colnames(theta_grid),
               function(a) sort(unique(as.numeric(theta_grid[, a])))),
        colnames(theta_grid))
    # A column holding one value across every cell is a fixed setting rather
    # than an axis of the grid: it contributes the same factor to every cell,
    # which cancels when the weights are normalised. Some are not quantities to
    # integrate at all, such as the `r = Inf` node a Poisson count grid carries
    # to pin the negative-binomial size, so they are dropped before an axis spec
    # is built for them.
    grids <- grids[vapply(grids, function(g) length(g) > 1L, logical(1))]
    if (length(grids) == 0L) return(NULL)
    specs <- .joint_axis_specs(grids, list(has_copy = TRUE), copy_slab = copy_slab,
                               folded_axes = folded_axes,
                               axis_refine = axis_refine)
    lapply(specs, function(sp) {
        if (.hp_is_logchol_col(.hyper_axis_bare(sp$name))) {
            sp$logchol_design <- logchol[[.hp_logchol_key(.hp_col_prefix(sp$name))]]
        }
        sp
    })
}

# The continuum levels of one axis: its distinct finite values, with the zero
# level dropped on a log-scale axis, where it is the point mass rather than a
# point of the continuum. The same convention `.hyper_axis_support()` applies
# before it measures a span.
.joint_axis_continuum_levels <- function(v, spec) {
    v <- sort(unique(as.numeric(v)))
    v <- v[is.finite(v)]
    if (isTRUE(spec$log_scale)) v <- v[v > 0]
    v
}

# The span each outer axis was integrated over, and the two terms that separate
# it from the nodes the caller wrote down.
#
# `axis_support` reports the interval the quadrature finally reached, which is
# the number a reader needs and not one they can decompose: it carries both the
# half node step the outermost cells own -- `k / (k - 1)` times the node range on
# the axis's integration coordinate, for `k` equally spaced nodes, so 2x at two
# nodes and 1.125x at nine -- and whatever refinement added past the end nodes.
# Reported side by side the two terms are readable off the fit:
#
#   nodes       range of the axis's continuum nodes, before any refinement pass
#   declared    support of that initial node set (the nodes plus the half step)
#   integrated  the region the final grid's measure integrates, read through
#               `.hyper_grid_supports()` with the final grid's `refining` tags,
#               so it is `axis_support`
#   refine      the mode that governed how far refinement could move the axis
#   n_nodes     continuum node count, initial and final
#
# An axis with fewer than two continuum levels has no span to report and is left
# out, matching what `.hyper_grid_supports()` does with it.
.joint_axis_span <- function(theta_grid_init, theta_grid_final, specs,
                             refining = NULL) {
    if (is.null(theta_grid_init) || is.null(specs)) return(NULL)
    theta_grid_init  <- as.matrix(theta_grid_init)
    theta_grid_final <- if (is.null(theta_grid_final)) theta_grid_init
                        else as.matrix(theta_grid_final)
    integrated <- .hyper_grid_supports(theta_grid_final, specs,
                                       refining = refining)
    out <- list()
    for (spec in specs) {
        a <- spec$name
        if (!a %in% colnames(theta_grid_init)) next
        lev0 <- .joint_axis_continuum_levels(theta_grid_init[, a], spec)
        if (length(lev0) < 2L) next
        lev1 <- if (a %in% colnames(theta_grid_final))
            .joint_axis_continuum_levels(theta_grid_final[, a], spec) else lev0
        mode <- if (!isTRUE(spec$refinable)) "none"
                else if (.hyper_axis_may_extend(spec)) "extend" else "densify"
        out[[a]] <- list(
            nodes      = range(lev0),
            declared   = .hyper_axis_support(theta_grid_init[, a], spec),
            integrated = integrated[[a]],
            refine     = mode,
            n_nodes    = c(initial = length(lev0), final = length(lev1))
        )
    }
    if (!length(out)) return(NULL)
    out
}

# The joint drivers' outer-grid refinement, run once the base grid is solved:
# the boundary / interior pass (opt in, `adaptive_grid`) and then the var-of-means
# consistency pass, both driven by the axis specs and a `kernel_fn` closure over
# the driver's own kernel call. Both drivers go through here so an axis refines
# the same way whichever door its fit entered by; each supplies what is its own --
# the specs, `kernel_fn`, `hp_fn` -- and reads the merged grid back. `axis_modes`
# is the outer mode a placement refit was laid from (`.nl_outer_mode_axes()`),
# which the consistency pass lays a collapsed axis's points at.
#
# Returns the merged `theta_grid`, `log_marginal`, `extras` and `refining_axis`,
# the two passes' records, and the number of nodes either added.
.joint_refine_outer_grid <- function(theta_grid, log_marginal, extras, specs,
                                     kernel_fn, hp_fn, adaptive_grid = FALSE,
                                     edge_thresh = 0.02, max_passes = 1L,
                                     consistency = TRUE, axis_modes = NULL) {
    refining_axis <- rep("", length(log_marginal))
    adaptive_info <- NULL
    consistency_info <- NULL
    n_added <- 0L
    if (isTRUE(adaptive_grid)) {
        refined <- .hyper_adaptive_refine_pass(
            theta_grid    = theta_grid,
            log_marginal  = log_marginal,
            extras        = extras,
            refining_axis = refining_axis,
            specs         = specs,
            kernel_fn     = kernel_fn,
            edge_thresh   = edge_thresh,
            max_passes    = max_passes,
            hp_fn         = hp_fn
        )
        n_added       <- n_added + length(refined$log_marginal) -
                         length(log_marginal)
        theta_grid    <- refined$theta_grid
        log_marginal  <- refined$log_marginal
        extras        <- refined$extras
        refining_axis <- refined$refining_axis
        adaptive_info <- refined$info
    }
    if (isTRUE(consistency)) {
        cons <- .hyper_consistency_pass(
            theta_grid    = theta_grid,
            log_marginal  = log_marginal,
            extras        = extras,
            refining_axis = refining_axis,
            specs         = specs,
            kernel_fn     = kernel_fn,
            hp_fn         = hp_fn,
            axis_modes    = axis_modes
        )
        if (cons$n_added > 0L) {
            theta_grid    <- cons$theta_grid
            log_marginal  <- cons$log_marginal
            extras        <- cons$extras
            refining_axis <- cons$refining_axis
            n_added       <- n_added + cons$n_added
        }
        consistency_info <- cons$info
    }
    list(theta_grid = theta_grid, log_marginal = log_marginal, extras = extras,
         refining_axis = refining_axis, adaptive_info = adaptive_info,
         consistency_info = consistency_info, n_added = n_added)
}

# Which cells of a per-cell `extras` list are solved and converged, solved and
# not, or never solved. A cell the screen dropped has no solve to judge; a cell
# whose log marginal is not finite failed outright and is not a stalled solve.
# `flagged` is FALSE when no cell carries a convergence flag, which is a kernel
# that reports none, not a grid of stalled cells.
#
# A cell whose Hessian at the returned point is not positive definite stopped
# at a point that is no mode, whatever its stopping rule said, so it is
# counted as stalled with the cells that never met the rule.
.joint_cell_convergence <- function(log_marginal, extras) {
    n <- length(log_marginal)
    has <- vapply(extras, function(e) !is.null(e$converged), logical(1))
    conv <- vapply(extras, function(e) isTRUE(e$converged) &&
                       !isTRUE(e$hessian_indefinite), logical(1))
    pruned <- vapply(extras, function(e) isTRUE(e$pruned), logical(1))
    solved <- !pruned & is.finite(log_marginal)
    list(flagged = length(extras) == n && any(has),
         converged = solved & conv, stalled = solved & !conv)
}

# Re-solve the cells whose inner Newton stopped without reaching a mode, each
# started from the mode of its nearest converged cell on the grid rather than
# from the broadcast start it was first given. The refinement passes start every
# cell of a level from the one heaviest base cell, which for a level laid across
# many rows is the right row for one of them; a converged neighbour one node
# away is where a linearly converging solve gains the most. A cell that still
# has no mode afterwards keeps its stalled record and is reported by
# `.nl_record_unconverged()`.
#
# `kernel_fn` is the drivers' refinement closure, which takes a per-cell start
# matrix (`x_init_per_cell`). Returns the merged `log_marginal` and `extras`
# and the count of cells re-solved and of those that converged.
.joint_resolve_unconverged <- function(theta_grid, log_marginal, extras,
                                       kernel_fn, hp_fn = NULL) {
    info <- list(n_stalled = 0L, n_resolved = 0L)
    out <- list(log_marginal = log_marginal, extras = extras, info = info)
    if (is.null(extras)) return(out)
    cc <- .joint_cell_convergence(log_marginal, extras)
    if (!cc$flagged) return(out)
    bad <- which(cc$stalled)
    info$n_stalled <- length(bad)
    out$info <- info
    good <- which(cc$converged & vapply(extras, function(e)
        !is.null(e$mode), logical(1)))
    if (!length(bad) || !length(good)) return(out)
    ref <- list(theta_grid = theta_grid[good, , drop = FALSE],
                modes = do.call(rbind, lapply(extras[good], function(e)
                    as.numeric(e$mode))))
    cells <- theta_grid[bad, , drop = FALSE]
    x0 <- .joint_nearest_grid_mode(cells, ref)
    fit_out <- kernel_fn(cells, warm_start = NULL, store_extras = TRUE,
                         x_init_per_cell = x0)
    new_lm <- as.numeric(fit_out$log_marginal)
    if (!is.null(hp_fn)) {
        hp_new <- hp_fn(cells)
        if (!is.null(hp_new) && length(hp_new) == length(bad)) {
            new_lm <- new_lm + hp_new
        }
    }
    new_ex <- fit_out$extras %||% vector("list", length(bad))
    for (k in seq_along(bad)) {
        e <- new_ex[[k]]
        if (!isTRUE(e$converged) || isTRUE(e$hessian_indefinite) ||
            !is.finite(new_lm[k])) next
        log_marginal[bad[k]] <- new_lm[k]
        extras[[bad[k]]] <- e
        info$n_resolved <- info$n_resolved + 1L
    }
    list(log_marginal = log_marginal, extras = extras, info = info)
}

# Record the cells whose inner solve never reached a mode, once the grid and
# its cell measure are final and before the weights are taken. A log marginal
# read off an iterate that is not a mode is not a Laplace approximation at a
# mode; how far it is from one is the cell's `score_max`, and a node taken out
# of a row leaves the row's quadratic join to bridge the gap, which moves the
# row's quantiles by more than a nearly converged iterate does. So the cells
# stay in the measure at the values they stopped at, and the
# fit says so: `nonconverged_cells` (an integer vector, empty when every solved
# cell converged, so a reader can tell "none" from "not recorded") and
# `nonconverged_mass`, the share of the posterior those cells carry, read
# against the same measure the weights use. The warning past the screen's mass
# gate is `.nl_warn_unconverged()`, said by the front door on the fit it
# returns: a placement refit replaces the grid this records, and a warning
# about the replaced grid names cells the caller never receives.
#
# A cell whose Hessian at the returned point is still not positive definite
# after that re-solve is a different case and is taken OUT of the measure. Its
# log marginal was formed from the determinant of a PD-enforced factor at a
# point that is no mode, so it is not a Laplace approximation of anything, and
# its stored precision cannot be factored, so every posterior draw that lands
# in it fails. Keeping its weight would make the reported moments integrate a
# cell the draws cannot. Recorded as `indefinite_cells` and `indefinite_mass`
# (the share it would have carried) and said by `.nl_warn_indefinite()`. When
# every solved cell is indefinite there is nothing to integrate instead, and
# the cells stay in the measure with the record saying so.
.nl_record_unconverged <- function(res) {
    lm <- as.numeric(res$log_marginal)
    n <- length(lm)
    res$nonconverged_cells <- integer(0)
    res$nonconverged_mass  <- 0
    res$indefinite_cells   <- integer(0)
    res$indefinite_mass    <- 0
    pruned <- as.logical(res$prune_mask %||% rep(FALSE, n))
    pruned[is.na(pruned)] <- FALSE
    indef <- as.logical(res$hessian_indefinite %||% logical(0))
    if (length(indef) == n) {
        ind <- which(is.finite(lm) & !pruned & indef %in% TRUE)
        if (length(ind)) {
            w0 <- .nl_normalise_weights_safe(lm, "outer grid",
                                             log_quad = res$log_quad)
            res$indefinite_cells <- ind
            res$indefinite_mass  <- sum(w0[ind], na.rm = TRUE)
            if (length(ind) < sum(is.finite(lm) & !pruned)) {
                lm[ind] <- -Inf
                res$log_marginal <- lm
            } else {
                res$indefinite_kept <- TRUE
            }
        }
    }
    conv <- res$converged
    if (is.null(conv) || length(conv) != n) return(res)
    bad <- which(is.finite(lm) & !pruned & !(as.logical(conv) %in% TRUE))
    if (!length(bad)) return(res)
    w0 <- .nl_normalise_weights_safe(lm, "outer grid", log_quad = res$log_quad)
    mass <- sum(w0[bad], na.rm = TRUE)
    res$nonconverged_cells <- bad
    res$nonconverged_mass  <- mass
    res
}

# Past the screen's mass gate, that much of the posterior sits on cells without
# a mode, and the fit says so once. `fn` names the front door.
.nl_warn_unconverged <- function(res, fn, gate_mass = .nl_screen("gate_mass")) {
    bad  <- res$nonconverged_cells
    mass <- res$nonconverged_mass
    if (length(bad) && is.finite(mass %||% NA_real_) && mass > gate_mass) {
        warning(sprintf(paste0(
            "%s: %d outer-grid cell%s whose inner Newton stopped at ",
            "control$max_iter without reaching a mode carr%s %s of the ",
            "posterior mass at the values they stopped at. Raise ",
            "control$max_iter or loosen control$tol; see `nonconverged_cells`."),
            fn, length(bad), if (length(bad) == 1L) "" else "s",
            if (length(bad) == 1L) "ies" else "y",
            format(signif(mass, 3))), call. = FALSE)
    }
    invisible(res)
}

# Cells left out of the integration because their Hessian stayed indefinite,
# said once by the front door on the fit it returns. Any such cell is said, not
# only one past a mass gate: a cell removed from the measure changes the
# posterior by the mass it carried, and the record is how a reader sees how
# much.
.nl_warn_indefinite <- function(res, fn) {
    bad  <- res$indefinite_cells
    if (!length(bad)) return(invisible(res))
    mass <- res$indefinite_mass
    warning(sprintf(paste0(
        "%s: %d outer-grid cell%s whose inner Hessian was not positive ",
        "definite at the point the inner Newton returned, even after a ",
        "re-solve from a converged neighbour, %s; %s %s of the posterior mass. ",
        "See `indefinite_cells` and `indefinite_mass`."),
        fn, length(bad), if (length(bad) == 1L) "" else "s",
        if (isTRUE(res$indefinite_kept))
            "were kept because no other cell was solved"
        else "were left out of the integration",
        if (length(bad) == 1L) "it carried" else "they carried",
        format(signif(mass, 3))), call. = FALSE)
    invisible(res)
}

# Convert a generic `new_cells` matrix [n_new x n_axes] back to the joint
# kernel's paired-vector `grids` representation. When `cp$has_copy = FALSE`
# the alpha entry stays `numeric(0)` (the no-copy contract the backend expects).
.joint_grids_from_cells <- function(new_cells, cp) {
    axes <- colnames(new_cells)
    out <- stats::setNames(lapply(axes, function(a)
        as.numeric(new_cells[, a])), axes)
    if (!cp$has_copy) out$alpha <- numeric(0)
    out
}

# Build the generic `kernel_fn(new_cells, warm_start, store_extras)` closure
# refinement passes around `backend$call_kernel`. Packs the joint kernel's
# per-cell side data (`.JOINT_CELL_FIELDS`) into a list of per-cell `extras` so
# the generic helpers can carry them along across refinement appends (and the
# warm-start chain still reads `extras[[idx0]]$mode`).
.joint_make_kernel_fn <- function(arms, prior, cp, backend, max_iter, tol,
                                  n_threads, x_init_default, store_Q,
                                  arm_names, cell_coupling = "separable",
                                  hessian_pd_mode = 0L,
                                  step_curvature_mode = 0L,
                                  inner_refresh = 1L,
                                  fixed_block_p = 0L,
                                  fixed_block_constraints = NULL,
                                  screen_tol = 0, screen_hp = NULL) {
    function(new_cells, warm_start = NULL, store_extras = FALSE,
             max_iter_override = NULL, n_threads_outer = 1L,
             inner_refresh_override = NULL, tol_override = NULL,
             x_init_per_cell = NULL,
             # Inner-Laplace skewness diagnostic, opt-in
             # like `store_extras`. Off by default so every existing caller
             # (adaptive-grid refinement, the outer Pareto-k re-evaluation)
             # is unaffected; .nlj_inner_skew_at_theta() is the only caller
             # that sets it, at a single (MAP) row of `new_cells`.
             compute_skew = FALSE, skew_idx = NULL,
             # Subspace debias: the kernel-facing request
             # list. Off by default, so every existing caller re-solves exactly
             # what it did; `.nl_subspace_debias_attach()` is the only caller
             # that sets it, over the fit's whole settled grid.
             debias = NULL,
             # Corrected integrated Laplace: the
             # kernel-facing request list, off by default for the same reason.
             cila = NULL,
             # A refinement round's cells screened against the grid they join
             # (`.hyper_refine_screen()`), at the fit's own tolerance.
             screen = NULL) {
        new_grids <- .joint_grids_from_cells(new_cells, cp)
        scr <- .nl_refine_screen_offset(new_cells, screen, screen_tol, screen_hp)
        slice_x_init <- if (!is.null(warm_start) && !is.null(warm_start$mode))
                        as.numeric(warm_start$mode) else x_init_default
        x_init_per_cell <- .joint_warm_start_per_cell(warm_start, x_init_per_cell,
                                                      nrow(new_cells))
        mi <- if (is.null(max_iter_override)) max_iter
              else as.integer(max_iter_override)
        # Per-call Shamanskii reuse override for the outer Pareto-k diagnostic
        #: the diagnostic re-solves only need the converged
        # log-marginal, so they run with factor reuse (grad-only scatter on the
        # off-factor steps) even when the fit itself keeps refresh = 1.
        ir <- if (is.null(inner_refresh_override)) inner_refresh
              else as.integer(inner_refresh_override)
        # Per-call inner-tol override for the diagnostic:
        # loosened to .K_DIAG_TOL since the Laplace log-marginal error is
        # O(tol^2). Never tighter than the fit's own tol.
        tl <- if (is.null(tol_override)) tol else max(as.numeric(tol_override), tol)
        # A call with n_threads_outer left at 1 solves its cells one after
        # another, each on the fit's whole thread grant (`n_threads`,
        # gcol33/tulpa#924); the refinement rounds and the outer Pareto-k
        # re-evaluation pass the fit's outer width, so their independent cells
        # run concurrently, each warm-started from the broadcast mode. Tiling is left off:
        # the IS batch is not a per-axis alpha lattice, so there is no tile
        # structure for the three-tier warm-start to exploit.
        # Both the precision and the per-cell fixed-effect block leave this
        # closure only inside `extras`, so a call that asks for no extras
        # requests neither. That keeps the outer Pareto-k batch -- hundreds of
        # cells, `store_extras` off -- from paying for either.
        res_x <- backend$call_kernel(arms, prior, cp, new_grids,
                                      mi, tl, n_threads,
                                      slice_x_init,
                                      isTRUE(store_Q) && isTRUE(store_extras),
                                      fixed_block_p = if (isTRUE(store_extras))
                                          as.integer(fixed_block_p) else 0L,
                                      fixed_block_constraints =
                                          fixed_block_constraints,
                                      arm_names = arm_names,
                                      n_threads_outer = as.integer(n_threads_outer),
                                      tile_warm = FALSE,
                                      cell_coupling = cell_coupling,
                                      hessian_pd_mode = hessian_pd_mode,
                                      step_curvature_mode = step_curvature_mode,
                                      inner_refresh = ir,
                                      x_init_per_cell = x_init_per_cell,
                                      compute_skew = compute_skew,
                                      skew_idx = skew_idx,
                                      debias = debias,
                                      cila = cila,
                                      prune_tol = if (is.null(scr)) 0 else screen_tol,
                                      screen_log_offset = scr)
        extras <- NULL
        if (isTRUE(store_extras)) {
            extras <- .joint_extras_from_res(res_x, nrow(new_cells))
        }
        list(log_marginal = res_x$log_marginal, extras = extras,
             modes = res_x$modes,
             inner_skew = res_x$inner_skew,
             inner_skew_gamma1 = res_x$inner_skew_gamma1,
             inner_skew_gamma1_declined = res_x$inner_skew_gamma1_declined,
             inner_skew_idx = res_x$inner_skew_idx,
             inner_skew_dropped = res_x$inner_skew_dropped,
             inner_skew_declined = res_x$inner_skew_declined,
             inner_skew_arms_declined = res_x$inner_skew_arms_declined,
             inner_is_z = res_x$inner_is_z,
             inner_is_sigma = res_x$inner_is_sigma,
             inner_is_log_joint = res_x$inner_is_log_joint,
             inner_is_declined = res_x$inner_is_declined,
             debias_draws_per_grid = res_x$debias_draws_per_grid,
             debias_idx = res_x$debias_idx,
             debias_accept = res_x$debias_accept,
             debias_declined = res_x$debias_declined,
             cila_log_w_per_grid = res_x$cila_log_w_per_grid,
             cila_fixed_per_grid = res_x$cila_fixed_per_grid,
             cila_log_marginal = res_x$cila_log_marginal,
             cila_variant = res_x$cila_variant,
             cila_declined = res_x$cila_declined,
             cila_fallback = res_x$cila_fallback)
    }
}

# The per-cell start matrix a kernel call hands its driver: the one it was
# given, else the `modes` a refinement round laid in `warm_start`
# (`.hyper_row_warm_starts()`) when it has a row per cell of the call.
.joint_warm_start_per_cell <- function(warm_start, x_init_per_cell, n_cells) {
    if (!is.null(x_init_per_cell)) return(x_init_per_cell)
    m <- warm_start$modes
    if (is.matrix(m) && nrow(m) == n_cells && all(is.finite(m))) m else NULL
}

# The per-cell side data a joint kernel result carries: one entry per cell of
# the outer grid, in the grid's own cell order. Refinement extends the grid cell
# by cell, so every one of these has to travel out through `extras` and back --
# a list left behind indexes the grid the fit had BEFORE refinement. The three
# passes that move them (slice the initial result, slice a refinement pass's
# result, glue the merged list back on) read this one table, so a field added to
# the kernel's output is added once rather than at three sites that then drift.
#
# `kind` is how a cell is taken out of a result and put back:
#   "row" -- an [n_grid x m] matrix; cell k is row k, a cell without one is NA
#   "int" -- a length-n_grid integer vector
#   "num" -- a length-n_grid numeric vector; a cell without one reads NA
#   "lgl" -- a length-n_grid logical vector; a cell without one reads FALSE
#   "elt" -- a length-n_grid list; cell k is element k
.JOINT_CELL_FIELDS <- list(
    # The cells the cheap-pass screen dropped unsolved, and each screened
    # cell's cheap log-marginal on the kernel's own scale. A refinement round
    # is screened against the grid its cells join (gcol33/tulpa#948); a cell
    # no screen saw reads FALSE / NA.
    list(res = "prune_mask",        extra = "pruned",    kind = "lgl"),
    list(res = "prune_cheap_log_marginal", extra = "cheap_lm", kind = "num"),
    list(res = "prune_screen_decrement", extra = "cheap_dec", kind = "num"),
    list(res = "modes",             extra = "mode",      kind = "row"),
    # The per-cell linear predictor and its within-cell variance: the pair a
    # grid-mixture predictive read draws a replicate from
    # (`.tulpa_eta_draws_grid()`, gcol33/tulpa#850). Present on a single-arm fit
    # only, the variance additionally only under `control$fitted_var`.
    list(res = "fitted_eta",        extra = "eta",       kind = "row"),
    list(res = "fitted_eta_var",    extra = "eta_var",   kind = "row"),
    list(res = "n_iter",            extra = "n_iter",    kind = "int"),
    # Whether the cell's inner Newton reached a mode. A refined cell's flag rides
    # here like every other per-cell field, so `converged` stays one entry per
    # cell of the merged grid; a shorter vector reads as all converged
    # (`.nested_converged_cells()`).
    list(res = "converged",         extra = "converged", kind = "lgl"),
    # The residual the inner Newton stopped at: the largest absolute component
    # of the joint penalized score at the cell's reported mode. `converged`
    # says whether the stopping rule fired; this says how far from stationary
    # a cell that stopped short is.
    list(res = "score_max",         extra = "score_max", kind = "num"),
    # Where the final-pass sum-to-zero log-determinant was read: off the
    # PD-enforced factor (`s2z_log_det_fallback`), and whether the field block
    # was indefinite there (`s2z_field_indefinite`). Counted by
    # `.tulpa_s2z_flag_cells()` over the merged grid.
    list(res = "s2z_log_det_fallback", extra = "s2z_log_det_fallback", kind = "lgl"),
    list(res = "s2z_field_indefinite", extra = "s2z_field_indefinite", kind = "lgl"),
    # Whether the Hessian at the cell's returned point was not positive
    # definite, so the point is no mode and its precision cannot be factored.
    # Re-solved by `.joint_resolve_unconverged()`; a cell that stays indefinite
    # is left out of the integration by `.nl_record_unconverged()`.
    list(res = "hessian_indefinite", extra = "hessian_indefinite", kind = "lgl"),
    list(res = "Q_csc_p_per_grid",  extra = "Q_csc_p",   kind = "elt"),
    list(res = "Q_csc_i_per_grid",  extra = "Q_csc_i",   kind = "elt"),
    list(res = "Q_csc_x_per_grid",  extra = "Q_csc_x",   kind = "elt"),
    list(res = "cov_block_per_grid", extra = "cov_block", kind = "elt")
)

# Per-cell extras for the first `n` cells of a joint kernel result: the fields
# it actually carries, under their `extras` names. Built for the initial
# cartesian pass (`.joint_init_extras_from_res`) and for every refinement
# append (`.joint_make_kernel_fn`) through the same walk.
.joint_extras_from_res <- function(res, n) {
    if (n == 0L) return(vector("list", 0L))
    lapply(seq_len(n), function(k) {
        e <- list()
        for (f in .JOINT_CELL_FIELDS) {
            v <- res[[f$res]]
            if (is.null(v)) next
            e[[f$extra]] <- switch(f$kind,
                                   row = as.numeric(v[k, ]),
                                   int = as.integer(v[k]),
                                   num = as.numeric(v[k]),
                                   lgl = as.logical(v[k]),
                                   elt = v[[k]])
        }
        e
    })
}

# Build the initial per-cell extras list from the initial joint kernel result,
# matching what `.joint_make_kernel_fn` would have produced for the cartesian
# pass. Refinement extends this list; `.joint_glue_extras_to_res` puts the
# refined extras back into `res` once integration is done.
.joint_init_extras_from_res <- function(res) {
    .joint_extras_from_res(res, length(res$log_marginal))
}

# The first cell carrying `field`, or 0 when none does.
#
# Which per-cell side data a fit holds is a property of the FIT (`store_Q`, the
# kernel's own outputs), but which cells hold it is not: a cheap-pass pruned
# cell is never solved, so it has no mode, no precision and no covariance
# block, whatever the fit stored. Reading cell 1 to decide for the whole list
# reads a pruned cell whenever the screen dropped the grid's first corner.
.joint_extras_first_with <- function(extras, field) {
    for (k in seq_along(extras)) {
        if (!is.null(extras[[k]][[field]])) return(k)
    }
    0L
}

# Glue refined extras + log_marginal + refining_axis back into the joint
# kernel result. Downstream `.nl_posterior_moments` / `.nl_attach_axis_sd`
# read `res$modes`,
# `res$log_marginal`, `res$refining_axis` directly; this keeps them in sync
# after refinement without touching their implementations.
#
# Every per-cell list is rewritten from `extras`, which the refinement passes
# extend cell for cell with the grid. A list left behind indexes the grid the
# fit had BEFORE refinement, and the per-cell precisions are what
# `tulpa_posterior_draws()` builds its mixture from, so a stale one stops
# prediction on a fit that otherwise converged.
.joint_glue_extras_to_res <- function(res, theta_grid_matrix, log_marginal,
                                      extras, refining_axis) {
    res$log_marginal  <- log_marginal
    res$n_grid        <- nrow(theta_grid_matrix)
    res$refining_axis <- refining_axis
    if (is.null(extras) || length(extras) == 0L) return(res)
    for (f in .JOINT_CELL_FIELDS) {
        k0 <- .joint_extras_first_with(extras, f$extra)
        if (k0 == 0L) next
        res[[f$res]] <- switch(
            f$kind,
            row = {
                m <- length(extras[[k0]][[f$extra]])
                do.call(rbind, lapply(extras, function(e) {
                    v <- e[[f$extra]]
                    if (is.null(v)) rep(NA_real_, m) else as.numeric(v)
                }))
            },
            int = vapply(extras, function(e) {
                v <- e[[f$extra]]
                if (is.null(v)) NA_integer_ else as.integer(v)
            }, integer(1)),
            num = vapply(extras, function(e) {
                v <- e[[f$extra]]
                if (is.null(v)) NA_real_ else as.numeric(v)
            }, numeric(1)),
            lgl = vapply(extras, function(e) isTRUE(e[[f$extra]]), logical(1)),
            elt = lapply(extras, `[[`, f$extra))
    }
    res
}

#' Per-cell constrained inverse-precision sub-block for a joint fit
#'
#' Runs the same per-outer-grid-cell selected-inversion the joint
#' nested-Laplace driver's own fixed-effect retention uses internally
#' (`extract_inner_vcov_block_cell()`, see the note above
#' `.joint_attach_grid_fixed()`), for a caller-chosen index set. A consumer
#' package computing its own standard-error correction on a subset of a
#' joint fit's latent coordinates -- a different `idx` than whatever the fit
#' itself retained on `$cov_block_per_grid` -- reruns the extraction through
#' this door instead of reimplementing the sparse selected inversion.
#'
#' @param Q_p_per_grid,Q_i_per_grid,Q_x_per_grid Per-outer-grid-cell CSC
#'   storage of the latent precision `Q` (`fit$Q_csc_p_per_grid` /
#'   `_i_per_grid` / `_x_per_grid`), one list element per cell.
#' @param n_x Integer, the latent dimension `Q` is built over
#'   (`fit$Q_csc_n`).
#' @param idx Integer vector (1-based) of latent coordinates to extract the
#'   sub-block for.
#' @param n_dense Integer, the number of leading `idx` entries that are
#'   dense fixed effects rather than field coordinates (see Details).
#' @param A_cols_list A list of integer vectors, one per sum-to-zero
#'   constraint group, each the (1-based) latent coordinates that group
#'   averages to zero.
#' @param field_marginal Logical; if `TRUE` (default) and
#'   `n_dense < length(idx)`, only the field coordinates' marginal variances
#'   are computed (a single Takahashi pass) rather than their full dense
#'   sub-block.
#' @param n_threads Integer, grid cells processed concurrently.
#'
#' @return A list of length `n_grid` (one per outer-grid cell), each element
#'   either `NULL` (the cell stored no `Q`, or its sparse Cholesky failed)
#'   or a dense `length(idx) x length(idx)` matrix: the constrained
#'   covariance sub-block for `idx`, with the sum-to-zero constraints in
#'   `A_cols_list` conditioned out by the same kriging correction the fit's
#'   own posterior draws use.
#'
#' @details When a field is present (`n_dense < length(idx)`), the
#'   extraction takes a cheap selected-inversion path: the dense
#'   (fixed-effect) block and the fixed-effect x field cross terms are
#'   exact; the field x field block is either its marginal diagonal
#'   (`field_marginal = TRUE`) or left at zero (`FALSE`) -- never the full
#'   field x field covariance.
#' @export
#' @examples
#' # Two outer-grid cells of a latent precision over two fixed effects and a
#' # ten-node random-walk field (sum-to-zero), each stored as the lower-triangle
#' # CSC a joint fit keeps per cell.
#' make_Q <- function(tau) {
#'   Q <- matrix(0, 12, 12)
#'   Q[1:2, 1:2] <- diag(2, 2) + 0.1
#'   Q[3:12, 3:12] <- tau * crossprod(diff(diag(10))) + diag(0.3, 10)
#'   Q[1, 3:12] <- Q[3:12, 1] <- 0.1
#'   L <- Matrix::tril(methods::as(methods::as(Q, "generalMatrix"),
#'                                 "CsparseMatrix"))
#'   list(p = L@p, i = L@i, x = L@x)
#' }
#' cells <- list(make_Q(1), make_Q(4))
#' blocks <- tulpa_joint_inner_vcov_blocks(
#'   lapply(cells, `[[`, "p"), lapply(cells, `[[`, "i"), lapply(cells, `[[`, "x"),
#'   n_x = 12L, idx = 1:12, n_dense = 2L, A_cols_list = list(3:12))
#' blocks[[1]][1:2, 1:2]        # fixed-effect covariance in the first cell
#' diag(blocks[[2]])[3:12]      # field marginal variances in the second
tulpa_joint_inner_vcov_blocks <- function(Q_p_per_grid, Q_i_per_grid, Q_x_per_grid,
                                          n_x, idx, n_dense, A_cols_list,
                                          field_marginal = TRUE, n_threads = 1L) {
  cpp_joint_inner_vcov_blocks(
    Q_p_per_grid = Q_p_per_grid, Q_i_per_grid = Q_i_per_grid,
    Q_x_per_grid = Q_x_per_grid,
    n_x = as.integer(n_x), idx = as.integer(idx), n_dense = as.integer(n_dense),
    A_cols_list = lapply(A_cols_list, as.integer),
    field_marginal = isTRUE(field_marginal), n_threads = as.integer(n_threads)
  )
}

#' Outer-grid posterior covariance of a joint fit's latent coordinates
#'
#' The law-of-total-covariance mixture over the outer grid,
#' \deqn{V = \sum_k w_k [C_k + (m_k - \bar m)(m_k - \bar m)^T],}
#' with `C_k` the same per-cell constrained block
#' [tulpa_joint_inner_vcov_blocks()] returns, accumulated inside the engine
#' without holding one block per cell. Memory is one `length(idx)` square
#' result whatever the grid size, so this is the door for an `idx` that
#' carries a field.
#'
#' @inheritParams tulpa_joint_inner_vcov_blocks
#' @param weights Numeric, one non-negative weight per outer-grid cell
#'   (normalized here). A cell of zero weight is never extracted.
#' @param modes Optional `n_grid x length(idx)` matrix of the cell modes on
#'   `idx`. `NULL` returns the within-cell term alone.
#'
#' @return A list: `vcov`, the mixture covariance; `mean`, the weighted mean
#'   of `modes` (`NULL` without them); `dense_blocks`, one
#'   `n_dense x n_dense` block per cell (`NULL` where the cell was not
#'   extracted); `failed`, the cells of positive weight whose block is
#'   missing (no stored precision, or its Cholesky failed). Such a cell keeps
#'   its weight and contributes its between-cell term only. With
#'   `field_marginal = TRUE` the field x field part of the within term is its
#'   diagonal.
#' @export
#' @examples
#' # Two outer-grid cells of a latent precision over two fixed effects and a
#' # ten-node random-walk field (sum-to-zero), each stored as the lower-triangle
#' # CSC a joint fit keeps per cell, with the cells' modes and weights.
#' make_Q <- function(tau) {
#'   Q <- matrix(0, 12, 12)
#'   Q[1:2, 1:2] <- diag(2, 2) + 0.1
#'   Q[3:12, 3:12] <- tau * crossprod(diff(diag(10))) + diag(0.3, 10)
#'   Q[1, 3:12] <- Q[3:12, 1] <- 0.1
#'   L <- Matrix::tril(methods::as(methods::as(Q, "generalMatrix"),
#'                                 "CsparseMatrix"))
#'   list(p = L@p, i = L@i, x = L@x)
#' }
#' cells <- list(make_Q(1), make_Q(4))
#' set.seed(1)
#' modes <- matrix(rnorm(2 * 12, sd = 0.2), 2, 12)
#' mix <- tulpa_joint_inner_vcov_mixture(
#'   lapply(cells, `[[`, "p"), lapply(cells, `[[`, "i"), lapply(cells, `[[`, "x"),
#'   n_x = 12L, idx = 1:12, n_dense = 2L, A_cols_list = list(3:12),
#'   weights = c(0.7, 0.3), modes = modes)
#' mix$mean[1:2]
#' mix$vcov[1:2, 1:2]
tulpa_joint_inner_vcov_mixture <- function(Q_p_per_grid, Q_i_per_grid,
                                           Q_x_per_grid, n_x, idx, n_dense,
                                           A_cols_list, weights, modes = NULL,
                                           field_marginal = TRUE,
                                           n_threads = 1L) {
  if (!is.null(modes)) modes <- as.matrix(modes) + 0
  cpp_joint_inner_vcov_mixture(
    Q_p_per_grid = Q_p_per_grid, Q_i_per_grid = Q_i_per_grid,
    Q_x_per_grid = Q_x_per_grid,
    n_x = as.integer(n_x), idx = as.integer(idx), n_dense = as.integer(n_dense),
    A_cols_list = lapply(A_cols_list, as.integer),
    weights = as.numeric(weights), modes = modes,
    field_marginal = isTRUE(field_marginal), n_threads = as.integer(n_threads)
  )
}

# ----------------------------------------------------------------------------
# Per-cell fixed-effect mode + precision for a joint fit.
#
# `.nested_fixed_moments()` (R/methods_generic.R) marginalizes the fixed effects
# over the outer grid by the law of total variance, and reads one representation
# to do it: `$grid_modes[[k]]` and `$grid_hessians[[k]]`, the cell-k fixed-effect
# mode and marginal precision. `tulpa_nested_laplace()` fills that pair from the
# stored per-cell precision; without it every reported interval is NA. This is
# the joint counterpart, filling the SAME pair so both tiers reach the one
# marginalizer -- and it is called from both the single-block and the multi-block
# joint driver, which is why it lives here rather than in either of them.
#
# Both joint layouts stack every arm's coefficients as a contiguous prefix of the
# latent vector (`.joint_layout()` / `.joint_multi_layout()` both start
# `beta_start` at 0 and run the arms consecutively), so the fixed-effect block is
# latent indices `1:n_fixed` -- the same span `.joint_fixed_layout()` names.
#
# The per-cell block arrives on `$cov_block_per_grid`, extracted by the inner
# Newton loop inside each cell's own solve through the same
# `extract_inner_vcov_block_cell()` the R-level `cpp_joint_inner_vcov_blocks()`
# drives, so the block is the constrained covariance the fit's own posterior
# draws are generated from -- the field sum-to-zero groups go in with the kernel
# call. Nothing here holds a precision: the whole grid's worth of them is never
# resident, and what the fit carries home is O(n_fixed^2) per cell.
#
# A decline is recorded on `$grid_fixed_declined` rather than left as an absent
# field, so a fit that reports NA intervals says why it does.
.joint_attach_grid_fixed <- function(res, n_fixed) {
    decline <- function(reason) {
        res$grid_modes <- NULL
        res$grid_hessians <- NULL
        res$grid_fixed_declined <- reason
        res
    }
    p <- as.integer(n_fixed %||% 0L)
    if (length(p) != 1L || is.na(p) || p < 1L) return(decline("no_fixed_effects"))

    # Read before the retention checks below, because it is upstream of them: a
    # cell whose inner solve never reached a mode has no fixed-effect
    # covariance to hand over, so the retention would report the absence it
    # trips over first and name a step downstream of the real cause.
    if (!.nested_any_weighted_converged(res)) return(decline("not_converged"))

    V <- res$cov_block_per_grid
    if (is.null(V)) return(decline("block_not_extracted"))

    modes <- res$modes
    w     <- res$weights
    n_grid <- length(V)
    # Every per-cell array must describe the SAME grid. A refinement pass that
    # rewrote the grid without carrying the blocks along would otherwise be
    # marginalized against stale cells.
    if (!is.matrix(modes) || nrow(modes) != n_grid || ncol(modes) < p ||
        is.null(w) || length(w) != n_grid) {
        return(decline("grid_misaligned"))
    }

    grid_hessians <- vector("list", n_grid)
    grid_modes    <- vector("list", n_grid)
    for (k in seq_len(n_grid)) {
        Vk <- V[[k]]
        ok <- is.matrix(Vk) && nrow(Vk) == p && ncol(Vk) == p
        Hk <- if (!ok) NULL else tryCatch(.spd_inverse(Vk), error = function(e) NULL)
        # A cell carrying no usable block is fatal only if it carries weight;
        # a zero-weight cell contributes nothing and its slot is left empty for
        # `.nested_fixed_moments()` to skip. The slot is SKIPPED rather than
        # assigned NULL: `l[[k]] <- NULL` removes the element, which on a
        # trailing empty cell shortens the list below `n_grid` and takes the
        # whole marginalization down on the length check.
        if (is.null(Hk)) {
            if (is.finite(w[k]) && w[k] > 0) {
                return(decline("cell_block_unavailable"))
            }
            next
        }
        grid_hessians[[k]] <- Hk
        grid_modes[[k]]    <- as.numeric(modes[k, seq_len(p)])
    }

    # A retention holding no cell the weights put mass on cannot produce a
    # coefficient table, so it is a decline with a reason rather than a pair of
    # lists `.nested_fixed_moments()` silently returns NULL on.
    filled <- !vapply(grid_hessians, is.null, logical(1))
    if (!any(filled & is.finite(w) & w > 0)) {
        return(decline("no_weighted_cell_block"))
    }

    res$grid_hessians <- grid_hessians
    res$grid_modes    <- grid_modes
    res$grid_fixed_declined <- NA_character_
    res
}

# The one place either joint driver settles what it retains for the fixed-effect
# marginalization. The per-cell blocks are already on the result; this turns
# them into the `$grid_modes` / `$grid_hessians` pair the marginalizer reads,
# and drops the raw blocks afterwards.
.joint_finalize_grid_fixed <- function(res, n_fixed, keep_grid_hessians,
                                       keep_joint = FALSE) {
    res <- if (isTRUE(keep_grid_hessians)) {
        .joint_attach_grid_fixed(res, n_fixed)
    } else {
        res$grid_fixed_declined <- "not_requested"
        res
    }
    # Same assembler the single/multi-block driver uses, so the two paths
    # cannot disagree about which cell the closure's coupling came from.
    if (isTRUE(keep_joint)) res$H_joint <- .nl_modal_joint_precision(res)
    res$cov_block_per_grid <- NULL
    res
}

# Latent dimension above which `force_sparse = "auto"` selects the sparse
# backend. Below it the dense factorization is the cheaper of the two: the
# sparse path pays a symbolic-analysis and indirection overhead that only earns
# its keep once the field is large enough for the fill-in saving to dominate.
.FORCE_SPARSE_AUTO_NX <- 1000L

# Resolve `control$force_sparse` to the boolean the kernels take. TRUE / FALSE
# pass through; "auto" compares the threshold against the latent dimension the
# kernel will actually build, supplied by the caller as a thunk because the
# single-block and multi-block paths reach it through different layouts
# (`.joint_layout` vs `.joint_multi_layout`) at different points. A thunk that
# fails or yields no finite n_x falls back to dense: the sparse path is the
# specialized one, so an unknown size should not silently opt into it.
.resolve_force_sparse <- function(force_sparse, n_x_fn) {
  if (is.logical(force_sparse) && length(force_sparse) == 1L &&
      !is.na(force_sparse)) {
    return(force_sparse)
  }
  if (!identical(force_sparse, "auto")) {
    stop("`control$force_sparse` must be TRUE, FALSE, or \"auto\".",
         call. = FALSE)
  }
  n_x <- tryCatch(n_x_fn(), error = function(e) NULL)
  if (is.null(n_x) || length(n_x) != 1L || !is.finite(n_x)) return(FALSE)
  n_x > .FORCE_SPARSE_AUTO_NX
}

# Which factorization the DENSE inner joint Newton uses on the Hessian it has
# already assembled densely. `force_sparse` above chooses the driver -- dense
# assembly or sparse assembly; this chooses the solver inside the dense one, and
# the two are independent. "auto" is the latent-dimension threshold, and the two
# explicit settings are what drive one problem through both backends.
.INNER_FACTORIZATION <- c("auto", "sparse", "dense")

.resolve_inner_factorization <- function(x) {
  if (is.null(x)) return(0L)
  if (!is.character(x) || length(x) != 1L || !(x %in% .INNER_FACTORIZATION)) {
    stop("`control$inner_factorization` must be one of ",
         paste(dQuote(.INNER_FACTORIZATION, FALSE), collapse = ", "), ".",
         call. = FALSE)
  }
  switch(x, auto = 0L, sparse = 1L, dense = -1L)
}

# The observations a joint fit is scored on: every arm's response, named by
# arm, and their total count -- the number of likelihood terms the arms
# contribute together, which is what nobs() and logLik() report.
.joint_attach_observations <- function(res, responses) {
    res$y <- lapply(responses, function(r) r$y)
    res$N <- sum(lengths(res$y))
    res
}

# Compute per-arm latent offsets so callers can decode `modes` back into
# per-arm (beta, re) blocks plus the shared spatial block(s). For BYM2 the
# spatial block is two sub-blocks (phi, theta); for ICAR/CAR_proper it's
# just phi.
.joint_layout <- function(arms, n_spatial_units, n_spatial_blocks,
                          spatial_block_names) {
    n_arms <- length(arms)
    p_arm  <- vapply(arms, function(a) ncol(a$X),     integer(1))
    n_re   <- vapply(arms, function(a) a$n_re_groups, integer(1))

    beta_start <- integer(n_arms)
    cur <- 0L
    for (k in seq_len(n_arms)) {
        beta_start[k] <- cur
        cur <- cur + p_arm[k]
    }
    re_start <- integer(n_arms)
    for (k in seq_len(n_arms)) {
        re_start[k] <- cur
        cur <- cur + n_re[k]
    }
    spatial_starts <- integer(n_spatial_blocks)
    for (b in seq_len(n_spatial_blocks)) {
        spatial_starts[b] <- cur
        cur <- cur + n_spatial_units
    }
    n_x <- cur

    out <- list(
        n_arms     = n_arms,
        p          = p_arm,
        n_re       = n_re,
        beta_start = beta_start,
        re_start   = re_start,
        n_x        = n_x
    )
    for (b in seq_len(n_spatial_blocks)) {
        out[[spatial_block_names[b]]] <- spatial_starts[b]
    }
    out
}

# The one R -> compiled-kernel dispersion boundary for the joint tier. Every
# R-level `phi` in the package is the engine convention (gaussian / lognormal:
# the residual VARIANCE); the kernels parameterize those two families by the
# residual SD. An arm carries its own family, so the arm scalar, the per-arm
# outer-grid override (`phi_grid_per_arm`, a per-cell vector or, on the batched
# entry, an [n_grid x n_batch] matrix whose shape is kept) and the per-species
# batch matrix (`phi_batch`, [n_arms x n_batch]) all convert against the same
# family here rather than at each of the joint entries.
#' @keywords internal
.joint_phi_args_to_kernel <- function(args) {
    arms <- args$arms_list
    if (is.null(arms)) return(args)
    fams <- vapply(arms, function(a) as.character(a$family %||% ""), character(1))
    for (k in seq_along(arms)) {
        arms[[k]]$phi <- .phi_to_kernel(fams[k], as.numeric(arms[[k]]$phi %||% 1.0))
    }
    args$arms_list <- arms
    g <- args$phi_grid_per_arm
    if (!is.null(g)) {
        for (k in seq_along(g)) {
            if (!is.null(g[[k]])) {
                g[[k]][] <- .phi_to_kernel(fams[k], as.numeric(g[[k]]))
            }
        }
        args$phi_grid_per_arm <- g
    }
    pb <- args$phi_batch
    if (!is.null(pb)) {
        pb <- as.matrix(pb)
        for (k in seq_len(nrow(pb))) {
            pb[k, ] <- .phi_to_kernel(fams[k], as.numeric(pb[k, ]))
        }
        args$phi_batch <- pb
    }
    args
}