R/fit_spde_nested.R

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

# Nested-Laplace integration backends for fit_spde().
#
# Two methods are supported:
#   * "grid" -- rectangular grid in log(range) x log(sigma) around the
#     prior modes. Cheap, defensible, but burns most of its budget on
#     points that the data have already ruled out.
#   * "ccd"  -- Central composite design centered on the joint posterior
#     mode of (range, sigma), oriented by the local Hessian. Mirrors
#     R-INLA's hyperparameter integration recipe; for k = 2 it uses
#     1 + 2*2 + 2^2 = 9 design points instead of n_grid^2 = 25 (default).
#
# Both backends return the same shape so the user-facing fit_spde()
# return contract is unchanged across methods.

# Outer Pareto-k-hat for the SPDE (range, sigma) integration: importance-sample
# the joint hyperparameter posterior on the log scale against the Gaussian
# proposal (theta_hat, L_scale, both in (log_range, log_sigma) space) and
# PSIS-smooth. The target is the SAME posterior the quadrature weights use, the
# marginal plus the hyperprior carried to (log range, log sigma)
# (`.spde_log_hyperprior()`). Runs with the RNG restored so the fit is
# unperturbed.
#
# Scored through the shared candidate dispatch
# (R/outer_pareto_candidates.R). `u_grid` / `w` are the integration nodes, which
# the GRID method has and whose weighted moments its proposal IS -- so the grid
# mixture is offered there. The CCD method's proposal is the mode-find's own
# Gaussian, so it passes neither and the mixture declines rather than reading a
# design spacing as a grid resolution.
.spde_pareto_k <- function(theta_hat, L_scale, spde_log_marginal, sp, n_samples,
                           u_grid = NULL, w = NULL,
                           proposal_source = "mode_hessian",
                           tail_points = NULL, hyperprior = "proper") {
  lt <- function(U) {
    r <- exp(U[, 1L]); s <- exp(U[, 2L])
    lm <- tryCatch(spde_log_marginal(r, s)$log_marginal,
                   error = function(e) rep(-Inf, length(r)))
    if (length(lm) != length(r)) return(rep(-Inf, length(r)))
    lm + .spde_log_hyperprior(r, s, sp, hyperprior)
  }
  .k_tail_cap_warn(tail_points, n_samples)
  kd <- .with_preserved_seed(
    tryCatch(.k_dispatch_report(
      .k_cand_spec(lt = lt, u_hat = theta_hat,
                   Su = L_scale %*% t(L_scale), u_grid = u_grid, w = w,
                   proposal_source = proposal_source),
      n_samples, tail_points = tail_points,
      scope = "fit_spde (log range, log sigma)"),
      error = function(e) NULL)
  )
  if (is.null(kd)) {
    list(pareto_k = NA_real_, is_ess = NA_real_,
         proposal_source = NA_character_,
         declined = .k_decline_label(.k_decline("degenerate_proposal",
                                               "the scorer errored")))
  } else kd
}

# ---------------------------------------------------------------------
# Method = "grid": rectangular grid over the SPDE hyperparameters, on the
# nested-Laplace ABI (re_idx / n_re_groups / sigma_re).
# ---------------------------------------------------------------------
fit_spde_nested_grid <- function(spde_log_marginal, sp, n_grid, spatial,
                                 diagnose_k = TRUE,
                                 k_samples = .nl_diag("k_samples"),
                                 k_tail_points = NULL,
                                 hyperprior = "proper") {
  range_mode <- sp$prior_range[1]
  sigma_mode <- sp$prior_sigma[1]

  lo_m <- .nl_grid_par("spde_direct", "lo_mult")
  hi_m <- .nl_grid_par("spde_direct", "hi_mult")
  range_grid <- exp(seq(log(range_mode * lo_m), log(range_mode * hi_m),
                        length.out = n_grid))
  sigma_grid <- exp(seq(log(sigma_mode * lo_m), log(sigma_mode * hi_m),
                        length.out = n_grid))

  grid <- expand.grid(range = range_grid, sigma = sigma_grid)
  result <- spde_log_marginal(grid$range, grid$sigma)

  hp <- .spde_hyperprior_record(grid$range, grid$sigma, sp, hyperprior)
  log_post <- result$log_marginal + hp$lp

  if (!any(is.finite(log_post))) {
    stop("fit_spde() nested grid: every hyperparameter cell returned a ",
         "non-finite log-marginal. Check the mesh, the data, and the range / ",
         "sigma PC priors.", call. = FALSE)
  }
  best <- which.max(log_post)
  log_max <- max(log_post[is.finite(log_post)])
  weights <- exp(log_post - log_max)
  weights[!is.finite(weights)] <- 0
  weights <- weights / sum(weights)

  # Outer k-hat: no mode Hessian on the grid path, so fit the Gaussian proposal
  # to the grid posterior moments in (log_range, log_sigma).
  kd <- list(pareto_k = NA_real_, is_ess = NA_real_,
             declined = .k_decline_label(.k_decline("not_requested")))
  if (isTRUE(diagnose_k)) {
    u_grid <- cbind(log(grid$range), log(grid$sigma))
    u_hat  <- as.numeric(crossprod(weights, u_grid))
    cen    <- sweep(u_grid, 2L, u_hat)
    Su     <- crossprod(cen * weights, cen); Su <- (Su + t(Su)) / 2
    Lk <- tryCatch(t(chol(Su)), error = function(e) NULL)
    kd <- if (is.null(Lk)) {
      list(pareto_k = NA_real_, is_ess = NA_real_,
           declined = .k_decline_label(.k_decline("degenerate_proposal",
                                                  "grid-moment covariance not positive definite")))
    } else .spde_pareto_k(u_hat, Lk, spde_log_marginal, sp, k_samples,
                          u_grid = u_grid, w = weights,
                          proposal_source = "grid_moment",
                          tail_points = k_tail_points, hyperprior = hyperprior)
  }

  # Outer-grid collapse visibility. This span is
  # already centred on the PC-prior mode (mode * 0.3 to mode * 3), unlike
  # the other families' universal fixed ceiling, but a data mode far from
  # the prior's can still rail it -- and `method = "grid"` is an explicit
  # opt-in (fit_spde()'s default `control$method` is already `"ccd"`, the
  # mode-Hessian path the auto-recenter uses), so this attaches the diagnostic
  # rather than silently overriding the user's chosen method with a refit.
  regime <- .joint_attach_pareto_k_regime(
    list(theta_grid = cbind(range = grid$range, sigma = grid$sigma),
        weights = weights))

  # The evidence against each cell's absolute measure in (log range,
  # log sigma), over the axes whose normalised density `log_post` carries.
  tg <- cbind(range = grid$range, sigma = grid$sigma)
  ev <- .nl_attach_evidence(
    list(log_marginal = log_post, log_hyperprior_axes = hp$axes,
         log_hyperprior_declined = hp$declined), tg,
    .joint_axis_specs_from_grid(tg, folded_axes = hp$axes))

  list(
    mode = NULL,
    log_marginal = result$log_marginal,
    log_hyperprior = log_post - result$log_marginal,
    log_evidence = ev$log_evidence,
    log_evidence_declined = ev$log_evidence_declined,
    log_hyperprior_declined = if (length(hp$declined)) hp$declined,
    converged = all(result$n_iter > 0),
    spatial = spatial,
    pareto_k = kd$pareto_k,
    pareto_k_is_ess = kd$is_ess,
    pareto_k_declined = .k_reason_of(kd),
    pareto_k_proposal_source = kd$proposal_source %||% NA_character_,
    pareto_k_first_pass = kd$first_pass_k %||% NA_real_,
    # The GPD tail the shape was fitted on. A raised or capped tail was
    # invisible on this path (gcol33/tulpa#692).
    pareto_k_tail_points = kd$tail_points %||% NA_integer_,
    pareto_k_scope = "outer (range, sigma) Gaussian proposal",
    pareto_k_regime          = regime$pareto_k_regime,
    pareto_k_grid_edge_axes  = regime$pareto_k_grid_edge_axes,
    pareto_k_grid_edge_sides = regime$pareto_k_grid_edge_sides,
    nested = list(
      method     = "grid",
      range_grid = grid$range,
      sigma_grid = grid$sigma,
      weights    = weights,
      pareto_k   = kd$pareto_k,
      n_iter     = result$n_iter,
      best_idx   = best,
      range_mean = sum(weights * grid$range),
      sigma_mean = sum(weights * grid$sigma),
      range_best = grid$range[best],
      sigma_best = grid$sigma[best]
    )
  )
}

# ---------------------------------------------------------------------
# Method = "ccd": find the joint posterior mode of (log_range, log_sigma),
# build a CCD grid oriented by the local negative-Hessian inverse, and
# integrate. Refits at the best grid point so `mode` is non-NULL.
# ---------------------------------------------------------------------
fit_spde_nested_ccd <- function(spde_log_marginal,
                                fit_spde_single,
                                sp, spatial,
                                diagnose_k = TRUE,
                                k_samples = .nl_diag("k_samples"),
                                k_tail_points = NULL,
                                mode_find = .nl_mode_find_tuning("spde"),
                                hyperprior = "proper",
                                spde_log_marginal_grad = NULL) {
  range_mode <- sp$prior_range[1]
  sigma_mode <- sp$prior_sigma[1]

  # Outer mode-find on log scale, box-constrained around the prior modes
  # (log range / sigma in [prior * 1e-2, prior * 1e2]). Nelder-Mead drifted
  # to (~0, +Inf) on weakly-identified data; L-BFGS-B with bounds keeps the
  # optimiser inside the prior support and lets the post-fit sanity check
  # below catch the rare cases where the mode still lands at a corner.
  obj <- function(theta) {
    r <- exp(theta[1])
    s <- exp(theta[2])
    lm <- tryCatch(
      spde_log_marginal(r, s)$log_marginal,
      error = function(e) NA_real_
    )
    if (!is.finite(lm)) return(1e10)
    lp <- .spde_log_hyperprior(r, s, sp, hyperprior)
    if (!is.finite(lp)) return(1e10)
    -(lm + lp)
  }
  # `spde_log_marginal_grad(log_range, log_sigma)` is the ANALYTIC (implicit-
  # differentiation) gradient of the log-marginal term alone (gcol33/tulpa#809);
  # the hyperprior term is a closed-form density with no compiled Newton solve
  # behind it, so a plain central difference of it costs nothing next to the
  # inner Laplace solve the marginal gradient replaces. Any failure (a Q that is
  # not PD at this theta, an unconverged inner solve) falls back to a central
  # difference of `obj` itself at that point, so `gr` never disagrees with what
  # `optim()`'s own numerical gradient would have produced there.
  gr <- if (is.null(spde_log_marginal_grad)) {
    NULL
  } else {
    function(theta) {
      out <- tryCatch({
        dlm <- spde_log_marginal_grad(theta[1], theta[2])
        if (!all(is.finite(dlm))) stop("non-finite marginal gradient")
        r <- exp(theta[1]); s <- exp(theta[2])
        eps <- 1e-6
        lp0 <- .spde_log_hyperprior(r, s, sp, hyperprior)
        lp_r <- .spde_log_hyperprior(exp(theta[1] + eps), s, sp, hyperprior)
        lp_s <- .spde_log_hyperprior(r, exp(theta[2] + eps), sp, hyperprior)
        if (!all(is.finite(c(lp0, lp_r, lp_s)))) stop("non-finite hyperprior")
        dlp <- c((lp_r - lp0) / eps, (lp_s - lp0) / eps)
        -(dlm + dlp)
      }, error = function(e) NULL)
      if (!is.null(out)) return(out)
      fd_step <- 1e-4
      vapply(seq_along(theta), function(j) {
        h <- fd_step * pmax(abs(theta[j]), 1)
        tp <- theta; tp[j] <- tp[j] + h
        tm <- theta; tm[j] <- tm[j] - h
        (obj(tp) - obj(tm)) / (2 * h)
      }, numeric(1))
    }
  }
  init  <- c(log(range_mode), log(sigma_mode))
  lower <- init - log(100)
  upper <- init + log(100)
  op <- .nl_lbfgsb_mode_find(par = init, fn = obj, gr = gr,
                             lower = lower, upper = upper,
                             tuning = mode_find)

  bad_mode <- function(op) {
    if (is.null(op)) return(TRUE)
    if (op$convergence != 0) return(TRUE)
    if (any(!is.finite(op$par))) return(TRUE)
    # Mode pinned to the optimisation box -> data are uninformative for
    # this hyper; fall back to the rectangular grid.
    eps <- 1e-3
    any(abs(op$par - lower) < eps) || any(abs(op$par - upper) < eps)
  }

  if (bad_mode(op)) {
    warning("nested-Laplace mode-find did not produce a usable mode ",
            "(degenerate or hit prior bounds); falling back to the ",
            "rectangular grid.")
    return(fit_spde_nested_grid(spde_log_marginal, sp, n_grid = 5L, spatial,
                                diagnose_k = diagnose_k, k_samples = k_samples,
                                k_tail_points = k_tail_points,
                                hyperprior = hyperprior))
  }
  theta_hat <- op$par
  range_hat <- exp(theta_hat[1])
  sigma_hat <- exp(theta_hat[2])

  # Local Hessian via stats::optimHess (central differences). obj is the
  # negative log-posterior, so its Hessian at the mode is the precision of
  # the Gaussian approximation to p(theta | y); the Cholesky of its inverse
  # scales the CCD design.
  H <- stats::optimHess(par = theta_hat, fn = obj)
  prec <- 0.5 * (H + t(H))   # symmetrise before the conditioning check
  ev <- eigen(prec, symmetric = TRUE, only.values = TRUE)$values
  # Reject degenerate Hessians: if the smallest eigenvalue is essentially
  # zero relative to the largest, the posterior is too flat for a CCD
  # design to wrap. Quadratic-rule integration would put nodes at extreme
  # log-scale values where the SPDE Newton fails.
  if (any(!is.finite(ev)) ||
      max(abs(ev)) <= 0 ||
      min(ev) <= 1e-6 * max(abs(ev))) {
    warning("nested-Laplace Hessian is degenerate at the mode ",
            "(condition unfit for CCD); falling back to the rectangular grid.")
    return(fit_spde_nested_grid(spde_log_marginal, sp, n_grid = 5L, spatial,
                                diagnose_k = diagnose_k, k_samples = k_samples,
                                k_tail_points = k_tail_points,
                                hyperprior = hyperprior))
  }
  sigma_post <- tryCatch(.spd_inverse(prec), error = function(e) NULL)
  if (is.null(sigma_post)) {
    warning("nested-Laplace Hessian inverse failed; ",
            "falling back to the rectangular grid.")
    return(fit_spde_nested_grid(spde_log_marginal, sp, n_grid = 5L, spatial,
                                diagnose_k = diagnose_k, k_samples = k_samples,
                                k_tail_points = k_tail_points,
                                hyperprior = hyperprior))
  }
  L <- t(chol(sigma_post))   # sigma_post = L L^T

  # CCD design (k = 2 -> 1 + 4 + 4 = 9 points). Sphere radius sqrt(2)*1.1
  # places the factorial corners at +/- 1.1 per axis -- INLA's standardized
  # scaling (f0 = 1.1) -- and `ccd_weights()` gives the matching corrected
  # R-INLA design weights (single source of truth with re_cov; see ccd_grid.R).
  # Map z to physical (log_range, log_sigma), then exponentiate. Cap log-scale
  # excursions at the optimisation box so a CCD axial that exceeds the prior
  # support snaps back to the boundary instead of overflowing to 1e+200.
  ccd <- ccd_grid(k = 2L, f_0 = sqrt(2) * 1.1)
  dnode <- ccd_weights(ccd)
  theta_grid <- ccd_to_theta(ccd$z, theta_hat, L, log_scale = FALSE)
  theta_grid[, 1] <- pmin(pmax(theta_grid[, 1], lower[1]), upper[1])
  theta_grid[, 2] <- pmin(pmax(theta_grid[, 2], lower[2]), upper[2])
  range_grid <- exp(theta_grid[, 1])
  sigma_grid <- exp(theta_grid[, 2])

  result <- spde_log_marginal(range_grid, sigma_grid)
  hp <- .spde_hyperprior_record(range_grid, sigma_grid, sp, hyperprior)
  log_post <- result$log_marginal + hp$lp
  # CCD quadrature: node weight = design weight (Delta_k) times exp(log_post),
  # the INLA convention int ~ sum_k Delta_k pi(theta_k).
  log_max <- max(log_post)
  weights <- dnode * exp(log_post - log_max)
  weights <- weights / sum(weights)

  best <- which.max(log_post)

  # Refit at the joint posterior mode so the result has a usable mode
  # (the Laplace posterior at the marginal-mode hypers).
  fit_at_mode <- fit_spde_single(range_hat, sigma_hat)

  # Outer k-hat reuses the mode-find's own Gaussian (theta_hat on the log scale,
  # L = chol of the posterior covariance) as the importance proposal -- the same
  # Gaussian the CCD design is oriented by.
  kd <- if (isTRUE(diagnose_k)) {
    .spde_pareto_k(theta_hat, L, spde_log_marginal, sp, k_samples,
                   tail_points = k_tail_points, hyperprior = hyperprior)
  } else list(pareto_k = NA_real_, is_ess = NA_real_,
              declined = .k_decline_label(.k_decline("not_requested")))

  list(
    mode             = fit_at_mode$mode,
    beta             = fit_at_mode$beta,
    spatial_effects  = fit_at_mode$spatial_effects,
    log_det_Q        = fit_at_mode$log_det_Q,
    log_marginal     = result$log_marginal,
    # A CCD's design weights reproduce moments and carry no volume, so no
    # evidence is read off them.
    log_evidence     = NA_real_,
    log_evidence_declined = "moment_rule_design",
    log_hyperprior_declined = if (length(hp$declined)) hp$declined,
    converged        = all(result$n_iter > 0),
    spatial          = spatial,
    range            = range_hat,
    sigma            = sigma_hat,
    pareto_k         = kd$pareto_k,
    pareto_k_is_ess  = kd$is_ess,
    pareto_k_declined = .k_reason_of(kd),
    pareto_k_proposal_source = kd$proposal_source %||% NA_character_,
    pareto_k_first_pass = kd$first_pass_k %||% NA_real_,
    # The GPD tail the shape was fitted on. A raised or capped tail was
    # invisible on this path (gcol33/tulpa#692).
    pareto_k_tail_points = kd$tail_points %||% NA_integer_,
    pareto_k_scope   = "outer (range, sigma) Gaussian proposal",
    nested = list(
      method        = "ccd",
      range_grid    = range_grid,
      sigma_grid    = sigma_grid,
      weights       = weights,
      pareto_k      = kd$pareto_k,
      n_iter        = result$n_iter,
      kind          = ccd$kind,
      best_idx      = best,
      range_mode    = range_hat,
      sigma_mode    = sigma_hat,
      range_mean    = sum(weights * range_grid),
      sigma_mean    = sum(weights * sigma_grid),
      range_best    = range_grid[best],
      sigma_best    = sigma_grid[best],
      hessian_logsc = H,
      f_0           = ccd$f_0,
      n_points      = ccd$n_points
    )
  )
}