R/fit_spde.R

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

#' Fit a spatial model using SPDE Laplace approximation
#'
#' Fits a GLM with a Matern spatial field via the SPDE approach.
#' Uses CHOLMOD sparse solver with optional nested Laplace for
#' hyperparameter integration.
#'
#' @param y Integer response vector.
#' @param X Design matrix.
#' @param spatial A `tulpa_spatial` object from [spatial_spde()] or
#'   [spatial_spde_custom()].
#' @param family Distribution family: `"binomial"`, `"poisson"`,
#'   `"neg_binomial_2"`, or `"gaussian"` (continuous-field geostatistics, with
#'   `phi` the observation-noise standard deviation).
#' @param n_trials Integer vector of trial sizes (binomial only).
#' @param range Spatial range parameter. If NULL, uses nested Laplace to
#'   integrate over range and sigma.
#' @param sigma Marginal standard deviation. If NULL, uses nested Laplace.
#' @param nested_laplace Logical. If TRUE (default when range/sigma are NULL),
#'   use nested Laplace approximation over hyperparameters.
#' @template phi
#' @param offset Optional fixed additive term on the linear predictor
#'   (`eta = offset + X beta + A w`), length `length(y)`; `NULL` -> no offset.
#' @param re_idx,n_re_groups,sigma_re Optional single iid random-intercept
#'   `(1 | g)` term alongside the Matern field: `re_idx` is a length-`length(y)`
#'   1-based group index, `n_re_groups` the number of groups, and `sigma_re` the
#'   (conditioned) random-effect SD. The field and the RE block are Laplace-
#'   marginalised jointly. `n_re_groups = 0` (default) is no RE term. Not
#'   supported for a fractional-nu field.
#' @param mode Inference method (the method is an argument, not a parallel
#'   verb): `"laplace"` (default) is the nested-Laplace integration over
#'   `(range, sigma)` documented here; `"nuts"` delegates to [tulpa_nuts_spde()]
#'   for exact HMC over the field (and, unless both `range` and `sigma` are
#'   fixed, the Matern hyperparameters). `mode = "nuts"` does not support an
#'   `offset` or a random-effect term, and its sampler knobs pass via `control`
#'   (see [tulpa_nuts_spde()]); it returns that sampler's draws object.
#' @param hyperprior `"proper"` (default) or `"flat"`, the prior the nested
#'   integration puts on `(range, sigma)`. `"proper"` carries the spec's PC
#'   range and PC sigma priors. `"flat"` keeps only the anchors the spec was
#'   given explicitly (`prior_range` / `prior_sigma` passed to
#'   [spatial_spde()]); an axis whose anchor was defaulted carries no density,
#'   is named in `log_hyperprior_declined` as `"flat_hyperprior"`, and the
#'   evidence declines. Not read by `mode = "nuts"`, where `"flat"` errors.
#' @param control A named list of numerical / tuning knobs (statistical
#'   arguments stay in the signature above). Recognized entries:
#'   \itemize{
#'     \item `method`: hyperparameter integration backend when nested Laplace is
#'       active. `"ccd"` (default) uses a central composite design centered on
#'       the joint posterior mode of `(range, sigma)`, oriented by the local
#'       Hessian (9 design points instead of `n_grid^2`), folding the PC priors
#'       from `spatial$prior_range` / `spatial$prior_sigma` into the integrated
#'       marginal and falling back to `"grid"` if the surface is too flat for a
#'       Hessian-based design. `"grid"` uses a rectangular grid in
#'       `log(range) x log(sigma)` around the prior modes.
#'     \item `n_grid`: grid points per hyperparameter dimension for
#'       `method = "grid"` (ignored under `"ccd"`). Default 5.
#'     \item `diagnose_k`: if TRUE (default), compute the outer
#'       Pareto-\eqn{\hat{k}} accuracy diagnostic (`$pareto_k`) by importance
#'       sampling the joint `(range, sigma)` posterior on the log scale against
#'       the Gaussian proposal that orients the integration. See [tulpa_psis()].
#'     \item `k_samples`: importance draws for `diagnose_k`. Default 500, each
#'       one extra batched SPDE marginal evaluation. It is a precision knob:
#'       the GPD tail size is held at the fraction the default budget implies,
#'       so a larger budget sharpens the same k-hat rather than moving it to a
#'       deeper quantile of the weight distribution (gcol33/tulpa#631).
#'     \item `k_tail_points`: expert override for that tail size, in upper-tail
#'       order statistics. Silently capped at 20% of the draws, beyond which
#'       body ratios enter the tail and bias the shape.
#'     \item `mode_find`: tuning for the outer `(range, sigma)` mode-find under
#'       `method = "ccd"`, as `list(factr =, ndeps =, maxit =)`; supply any
#'       subset. `ndeps` is the central-difference step for `optim()`'s
#'       numerical gradient on the log scale (default 1e-2): it must clear the
#'       inner solver's convergence tolerance, and a step wide enough that its
#'       truncation error exceeds the reduction the line search chases near a
#'       flat optimum leaves L-BFGS-B aborting at the mode it just reached, in
#'       which case the CCD design declines to the rectangular grid. `factr` is
#'       the relative-reduction stop in units of `.Machine$double.eps` (default
#'       1e5); `maxit` the iteration cap (default 300).
#'     \item `max_iter`: maximum Newton iterations. Default 100.
#'     \item `tol`: Newton convergence tolerance. Default 1e-6.
#'     \item `n_threads`: OpenMP threads. Default 1.
#'     \item `checkpoint`: grid-cell checkpoint/resume spec
#'       `list(path =, resume =)`. Each solved `(range, sigma)` cell is appended
#'       to `path`; a `resume = TRUE` run loads the finished cells and re-solves
#'       only the rest, so a killed or rebooted fit resumes instead of
#'       restarting. `resume = FALSE` starts a fresh file. Default `NULL` (off).
#'   }
#'
#' @return A list with:
#'   \itemize{
#'     \item `mode`: mode of the latent field (beta + mesh node effects)
#'     \item `log_marginal`: log marginal likelihood
#'     \item `converged`: convergence flag
#'     \item `spatial`: the spatial specification (for prediction)
#'     \item `pareto_k`, `pareto_k_is_ess`: outer Pareto-\eqn{\hat{k}} and its
#'       importance-sampling ESS (`NA` when `diagnose_k = FALSE`)
#'     \item `pareto_k_proposal_source`, `pareto_k_first_pass`: which proposal
#'       family the reported k-hat came from (several candidates are scored and
#'       the best kept), and the k-hat of the first pass -- the proposal exactly
#'       as placed, before refinement. A large gap says the placement is poor
#'       even where the verdict is fine.
#'     \item `nested`: nested Laplace results (if used)
#'   }
#'
#' @references
#' Lindgren, Rue & Lindstrom (2011). An explicit link between Gaussian fields
#' and Gaussian Markov random fields: the stochastic partial differential
#' equation approach. \emph{JRSS-B} 73(4):423-498.
#' Rue, Martino & Chopin (2009). Approximate Bayesian inference for latent
#' Gaussian models by using integrated nested Laplace approximations.
#' \emph{JRSS-B} 71(2):319-392.
#' @examples
#' \donttest{
#' if (requireNamespace("fmesher", quietly = TRUE)) {
#'   set.seed(1)
#'   n <- 200L
#'   coords <- cbind(runif(n), runif(n))
#'   mesh <- fmesher::fm_mesh_2d(loc = coords, max.edge = c(0.15, 0.4), cutoff = 0.05)
#'   fem  <- fmesher::fm_fem(mesh)
#'   A    <- as(fmesher::fm_basis(mesh, loc = coords), "CsparseMatrix")
#'   spec <- spatial_spde_custom(C = fem$c0, G = fem$g1, A = A, nu = 1,
#'                               prior_range = c(0.3, 0.5), prior_sigma = c(0.6, 0.05))
#'   w <- as.numeric(rnorm(spec$n_mesh, 0, 0.6)); w <- w - mean(w)
#'   x <- rnorm(n)
#'   y <- rpois(n, exp(2.0 + 0.5 * x + as.numeric(spec$A %*% w)))
#'   fit <- fit_spde(y = y, X = cbind(1, x), spatial = spec, family = "poisson")
#'   fit$nested$range_mean
#' }
#' }
#' @export
fit_spde <- function(y, X, spatial,
                     family = "binomial", n_trials = NULL,
                     range = NULL, sigma = NULL,
                     nested_laplace = is.null(range) || is.null(sigma),
                     phi = 1.0, offset = NULL,
                     re_idx = NULL, n_re_groups = 0L, sigma_re = 1.0,
                     mode = c("laplace", "nuts"),
                     hyperprior = c("proper", "flat"),
                     control = list()) {

  mode <- match.arg(mode)
  family <- .canonical_family(family)
  .family_or_stop(family)
  hyperprior <- .hp_choice(match.arg(hyperprior))
  if (mode == "nuts" && identical(hyperprior, "flat")) {
    stop("`hyperprior = \"flat\"` is not read by mode = 'nuts', which samples ",
         "(range, sigma) under the spec's PC priors.", call. = FALSE)
  }
  if (!inherits(spatial, "tulpa_spatial") || spatial$type != "spde") {
    stop("spatial must be an SPDE tulpa_spatial object", call. = FALSE)
  }
  .check_spde_rows(spatial, length(y), "fit_spde()")
  if (mode == "nuts") {
    # The inference method is an argument, not a parallel verb: delegate to the
    # NUTS engine (tulpa_nuts_spde). The Laplace-path-only structure (an offset,
    # a random-effect block) is unsupported there.
    if (!is.null(offset)) {
      stop("`offset` is not supported on mode = 'nuts' (Laplace-path only).",
           call. = FALSE)
    }
    if (!is.null(re_idx) && n_re_groups > 0L) {
      stop("A random-effect term is not supported on mode = 'nuts' for an SPDE ",
           "field; use mode = 'laplace'.", call. = FALSE)
    }
    return(tulpa_nuts_spde(
      y = as.numeric(y), X = as.matrix(X), spatial = spatial, family = family,
      n_trials = n_trials,
      joint = is.null(range) || is.null(sigma),
      range = range, sigma = sigma,
      control = control))
  }
  tulpa_check_control(control, .CONTROL_KEYS$spde, "fit_spde")

  # Perf/numerical knobs live in `control = list()` (matching tulpa() /
  # tulpa_nested_laplace()); the signature carries only statistical arguments.
  method     <- match.arg(control$method %||% "ccd", c("ccd", "grid"))
  n_grid     <- as.integer(control$n_grid %||% 5L)
  diagnose_k <- isTRUE(control$diagnose_k %||% TRUE)
  k_samples  <- as.integer(control$k_samples %||% .nl_diag("k_samples"))
  k_tail_pts <- control$k_tail_points
  mode_find  <- .nl_mode_find_tuning("spde", control)
  max_iter   <- as.integer(control$max_iter %||% 100L)
  tol        <- control$tol %||% 1e-6
  n_threads  <- as.integer(control$n_threads %||% 1L)
  checkpoint <- control$checkpoint
  # The cheap screen, the per-row predictive-variance loop and the two
  # correction requests, resolved exactly as tulpa_nested_laplace() resolves
  # them. They were unreachable on this path until the SPDE entry went through
  # the shared bundle (gcol33/tulpa#699), which hardcoded prune_tol = 0.
  prune       <- isTRUE(control$prune %||% .nl_screen("prune")[["spde"]])
  prune_tol   <- .nl_prune_tol_from_control(
    control,
    .nl_check_prune_tol(control$prune_tol %||% .nl_screen("prune_tol")))
  prune_tol_eff <- if (prune) prune_tol else 0
  screen_it   <- .nl_check_screen_iters(
    control$screen_iters %||% .nl_screen("iters"))
  fitted_var  <- isTRUE(control$fitted_var %||% TRUE)
  debias_req  <- control$subspace_debias
  cila_req    <- control$cila

  if (!is.null(offset)) {
    offset <- as.numeric(offset)
    if (length(offset) != length(y)) {
      stop(sprintf("length(offset) (%d) must equal length(y) (%d).",
                   length(offset), length(y)), call. = FALSE)
    }
  }

  # Grid-cell checkpoint/resume. `checkpoint = list(path =,
  # resume =)` appends each solved (range, sigma) cell to `path`; a resume loads
  # finished cells and re-solves only the rest. The outer integrator
  # (CCD / grid) calls spde_log_marginal() many times across mode search and the
  # final grid -- every cell is keyed by its (range, sigma) coordinate, so all
  # those calls share one file and any previously solved coordinate is reused.
  .ckpt <- .nl_checkpoint_args(list(checkpoint = checkpoint))
  if (nzchar(.ckpt$path) && !isTRUE(.ckpt$resume) && file.exists(.ckpt$path)) {
    file.remove(.ckpt$path)
  }

  y <- as.numeric(y)
  n_obs <- length(y)
  X <- as.matrix(X)
  n_trials <- .validate_glm_design(y, X, n_trials, "fit_spde()")$n_trials

  sp <- spatial  # shorthand

  # v10 nested-Laplace ABI requires re_idx / n_re_groups / sigma_re even when
  # there is no formula-side RE term. A single iid random-intercept `(1 | g)`
  # can ride alongside the Matern field, conditioned on `sigma_re` (like the
  # areal spatial Laplace path); the whole latent block -- field + RE -- is
  # Laplace-marginalised jointly inside the kernel. When there is no RE term
  # these pin to the no-RE defaults so every nested-Laplace call is consistent.
  has_re <- !is.null(re_idx) && n_re_groups > 0L
  if (has_re) {
    re_idx <- as.integer(re_idx)
    if (length(re_idx) != n_obs) {
      stop(sprintf("length(re_idx) (%d) must equal length(y) (%d).",
                   length(re_idx), n_obs), call. = FALSE)
    }
    if (anyNA(re_idx) || min(re_idx) < 1L || max(re_idx) > n_re_groups) {
      stop(sprintf(paste0("`re_idx` must be 1-based integers in ",
                          "[1, n_re_groups = %d]; got range [%d, %d]."),
                   n_re_groups, min(re_idx), max(re_idx)), call. = FALSE)
    }
    if (!is.numeric(sigma_re) || length(sigma_re) != 1L ||
        !is.finite(sigma_re) || sigma_re <= 0) {
      stop("`sigma_re` must be a positive scalar.", call. = FALSE)
    }
  }
  no_re_idx       <- if (has_re) re_idx      else rep(0L, n_obs)
  no_re_n_groups  <- if (has_re) as.integer(n_re_groups) else 0L
  no_re_sigma     <- if (has_re) as.numeric(sigma_re)    else 1.0

  # Fractional nu integrates the operator-based rational SPDE:
  # each (range, sigma) cell assembles its own (Q, A_eff) via the validated R
  # oracle and is solved by the precomputed C++ fit, with 0.5 log|Q| folded into
  # the per-cell marginal (done inside .spde_laplace_fractional_at). The outer
  # CCD / grid integrators are reused unchanged -- they read only $log_marginal
  # and $n_iter. The cpp grid-cell checkpoint is integer-path only.
  is_frac   <- .spde_nu_is_fractional(sp$nu)
  order_rat <- sp$rational_order %||% 2L
  if (is_frac && has_re) {
    stop("A random-effect term alongside a fractional-nu SPDE field is not ",
         "supported yet; use an integer nu, or drop the (1 | g) term.",
         call. = FALSE)
  }
  if (is_frac && nzchar(.ckpt$path)) {
    warning("Grid-cell checkpointing is not supported for fractional nu; ",
            "the fractional SPDE grid runs without it.", call. = FALSE)
  }

  spde_log_marginal <- function(range_vec, sigma_vec) {
    if (is_frac) {
      K  <- length(range_vec)
      lm <- numeric(K); ni <- integer(K)
      for (k in seq_len(K)) {
        f <- .spde_nested_logmarginal_at(
          spatial = sp, range = range_vec[k], sigma = sigma_vec[k],
          y = y, X = X, family = family, phi = phi, n_trials = n_trials,
          order = order_rat, max_iter = max_iter, tol = tol,
          n_threads = n_threads, offset = offset
        )
        lm[k] <- f$log_marginal
        ni[k] <- f$n_iter
      }
      return(list(log_marginal = lm, n_iter = ni))
    }
    res <- cpp_nested_laplace_spde(
      y = y, n_trials = n_trials, X = X,
      re_idx = no_re_idx, n_re_groups = no_re_n_groups,
      sigma_re = no_re_sigma,
      A_x = sp$A_x, A_i = sp$A_i, A_p = sp$A_p,
      n_obs = n_obs, n_mesh = sp$n_mesh,
      C0_diag = sp$C0_diag,
      G1_x = sp$G1_x, G1_i = sp$G1_i, G1_p = sp$G1_p,
      range_grid = range_vec, sigma_grid = sigma_vec,
      nu = sp$nu,
      family = family, phi = .phi_to_kernel(family, phi),
      max_iter = max_iter, tol = tol, n_threads = n_threads,
      checkpoint_path = .ckpt$path,
      offset_nullable = offset,
      debias = debias_req, cila = cila_req,
      prune_tol = prune_tol_eff, screen_iters = screen_it,
      compute_fitted_var = fitted_var,
      # The grid and the CCD design weigh a cell by its marginal plus the PC
      # prior on (log range, log sigma), so a screened solve ranks with both.
      screen_log_offset = if (prune_tol_eff > 0)
                            .spde_log_hyperprior(range_vec, sigma_vec, sp,
                                                 hyperprior)
    )
    res
  }

  # Analytic (log range, log sigma) gradient of the log-marginal, for the CCD
  # outer mode-find's L-BFGS-B (fit_spde_nested_ccd() -> .nl_lbfgsb_mode_find()).
  # cpp_spde_laplace_gradient() implements exactly the alpha = 2 (nu = 1)
  # operator with no offset and no random-effect term (gcol33/tulpa#809);
  # outside that scope the mode-find keeps its central-difference gradient.
  spde_log_marginal_grad <- if (!is_frac && isTRUE(all.equal(sp$nu, 1)) &&
                                is.null(offset) && !has_re) {
    function(log_range, log_sigma) {
      g <- cpp_spde_laplace_gradient(
        y = y, n_trials = n_trials, X = X,
        A_x = sp$A_x, A_i = sp$A_i, A_p = sp$A_p,
        n_obs = n_obs, n_mesh = sp$n_mesh,
        C0_diag = sp$C0_diag,
        G1_x = sp$G1_x, G1_i = sp$G1_i, G1_p = sp$G1_p,
        log_range = log_range, log_sigma = log_sigma,
        nu = sp$nu, family = family, phi = .phi_to_kernel(family, phi),
        max_iter = max_iter, tol = tol, n_threads = n_threads
      )
      c(g$grad_log_range, g$grad_log_sigma)
    }
  } else NULL

  fit <- if (nested_laplace && (is.null(range) || is.null(sigma))) {
    if (method == "grid") {
      fit_spde_nested_grid(spde_log_marginal, sp, n_grid, spatial,
                           diagnose_k = diagnose_k, k_samples = k_samples,
                           k_tail_points = k_tail_pts,
                           hyperprior = hyperprior)
    } else {
      fit_spde_nested_ccd(spde_log_marginal,
                                 fit_spde_single = function(r, s) {
                                   laplace_spde_at(
                                     y = y, n_trials = n_trials, X = X,
                                     spatial = sp, family = family, phi = phi,
                                     range = r, sigma = s,
                                     max_iter = max_iter, tol = tol,
                                     n_threads = n_threads, offset = offset,
                                     re_idx = no_re_idx,
                                     n_re_groups = no_re_n_groups,
                                     sigma_re = no_re_sigma
                                   )
                                 },
                                 sp = sp, spatial = spatial,
                          diagnose_k = diagnose_k, k_samples = k_samples,
                          k_tail_points = k_tail_pts,
                          mode_find = mode_find,
                          hyperprior = hyperprior,
                          spde_log_marginal_grad = spde_log_marginal_grad)
    }
  } else {
    # --- Single-point Laplace at fixed hyperparameters ---
    # Delegate to the shared helper used by dispatch_laplace_spatial so the
    # SPDE Laplace call site stays singular.
    result <- laplace_spde_at(
      y = y, n_trials = n_trials, X = X, spatial = sp,
      family = family, phi = phi,
      range = range, sigma = sigma,
      max_iter = max_iter, tol = tol, n_threads = n_threads,
      offset = offset,
      re_idx = no_re_idx, n_re_groups = no_re_n_groups, sigma_re = no_re_sigma
    )

    list(
      mode = result$mode,
      beta = result$beta,
      spatial_effects = result$spatial_effects,
      log_det_Q = result$log_det_Q,
      log_marginal = result$log_marginal,
      n_iter = result$n_iter,
      converged = result$converged,
      spatial = spatial,
      range = result$range,
      sigma = result$sigma,
      nested = NULL
    )
  }

  # Prediction support (predict.tulpa_fit se.fit): the family and the
  # dispersion, in the one engine convention (gaussian / lognormal: the
  # residual variance), the same `$phi` a front-door tulpa() fit carries.
  fit$family <- fit$family %||% family
  fit$phi    <- phi

  # Fixed-effect mode + marginal curvature so the generic accessors (coef /
  # summary / confint / vcov / predict) work on a direct or auto-mode SPDE
  # fit. The nested grid path stores no per-cell modes, so anchor one refit
  # at the posterior-mean hyperparameters; H_beta is the field-marginal Schur
  # complement (.marginal_H_beta_spde), the same block mode = "laplace"
  # attaches.
  if (is.null(fit$mode) && !is.null(fit$nested)) {
    range_hy <- fit$nested$range_mean %||% fit$nested$range_best
    sigma_hy <- fit$nested$sigma_mean %||% fit$nested$sigma_best
    anchor <- laplace_spde_at(
      y = y, n_trials = n_trials, X = X, spatial = sp,
      family = family, phi = phi,
      range = range_hy, sigma = sigma_hy,
      max_iter = max_iter, tol = tol, n_threads = n_threads,
      offset = offset,
      re_idx = no_re_idx, n_re_groups = no_re_n_groups, sigma_re = no_re_sigma
    )
    fit$mode            <- anchor$mode
    fit$beta            <- fit$beta %||% anchor$beta
    fit$spatial_effects <- fit$spatial_effects %||% anchor$spatial_effects
    fit$range           <- fit$range %||% range_hy
    fit$sigma           <- fit$sigma %||% sigma_hy
  }
  if (!is.null(fit$mode)) {
    range_hy <- fit$range %||% fit$nested$range_mean %||% fit$nested$range_best
    sigma_hy <- fit$sigma %||% fit$nested$sigma_mean %||% fit$nested$sigma_best
    fit$H_beta <- fit$H_beta %||% tryCatch(
      .marginal_H_beta_spde(
        mode = fit$mode, X = X, spatial = sp,
        family = family, phi = phi,
        n_trials = n_trials, offset = offset,
        range_val = range_hy, sigma_val = sigma_hy,
        re_idx = no_re_idx, n_re_groups = no_re_n_groups, sigma_re = no_re_sigma,
        y = y
      ),
      error = function(e) NULL
    )
  }

  .finalize_fit(fit, backend = "spde",
                n_fixed = ncol(X), fixed_names = colnames(X),
                data = list(y = y, n_trials = n_trials, model_matrix = X,
                           family = family, offset = offset, phi = phi))
}


# The SPDE projector A maps observations to mesh nodes, one row per
# observation, and every SPDE kernel indexes it by the observation index with no
# bound of its own. A longer response than A has rows used to be accepted, the
# extra rows reaching the likelihood with no field term; a shorter one reached
# the kernel's CSC check as "A: row index 10 at slot 11 is outside [0, 10)"
# (gcol33/tulpa#908). One check, at every R door into an SPDE fit.
#' @keywords internal
.check_spde_rows <- function(spatial, n_obs, where) {
  n_a <- tryCatch(nrow(spatial$A), error = function(e) NULL)
  if (is.null(n_a) || n_a != n_obs) {
    stop(where, ": the SPDE projector matrix A has ", n_a %||% "?",
         " row(s) but the response has ", n_obs, " observation(s). Build the ",
         "SPDE spec from the same data (spatial_spde(~ lon + lat, data = ",
         "<data>)), one row per observation.", call. = FALSE)
  }
  invisible(TRUE)
}