R/nested_laplace_joint_backends.R

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

# Single-block backend dispatch table + kernel call.
#
# Part of the joint nested-Laplace driver; the public entry point
# tulpa_nested_laplace_joint() lives in nested_laplace_joint.R.

# --- backend dispatch table --------------------------------------------------
#
# Adding a new joint backend means appending one entry here. Each entry
# bundles the four backend-specific concerns: grid construction, kernel call,
# theta-grid materialisation for the result, and latent-layout metadata.

# Each entry owns three backend-specific concerns: grid construction,
# theta-grid materialisation, and latent-layout metadata. The kernel call is
# shared across all single-block backends -- they all route through
# `.joint_call_kernel_via_multi()` which packs the single-block prior into
# a length-1 multi-block spec and dispatches via
# `cpp_nested_laplace_joint_multi`. This is the J-E unification: one
# inner C++ entry, one R-side post-processing path. The legacy
# per-backend `cpp_nested_laplace_joint_{bym2,icar,car_proper}` shims and
# their bespoke LatentBlock construction are gone.

# The one `call_kernel` every single-block backend gets. The backend name is
# the only thing that varies across them, so the forwarding signature is
# written once here rather than per entry; `.joint_backends` reads
# `call_kernel` polymorphically (see `nested_laplace_joint_helpers.R`, which
# holds the entry but not its name), so the closure stays on the entry.
.make_joint_kernel_caller <- function(backend_name) {
    force(backend_name)
    function(arms, prior, cp, grids, max_iter, tol,
             n_threads, x_init, store_Q = FALSE,
             arm_names = NULL, n_threads_outer = 1L,
             tile_warm = TRUE, prune_tol = 0.0,
             force_sparse = FALSE,
             cell_coupling = "separable",
             hessian_pd_mode = 0L,
             step_curvature_mode = 0L,
             inner_refresh = 1L,
             x_init_per_cell = NULL,
             compute_skew = FALSE,
             skew_idx = NULL,
             fixed_block_p = 0L,
             fixed_block_constraints = NULL,
             debias = NULL,
             cila = NULL,
             screen_log_offset = NULL,
             screen_only = FALSE) {
        .joint_call_kernel_via_multi(backend_name, arms, prior, cp, grids,
                                      max_iter, tol, n_threads,
                                      x_init, store_Q, arm_names,
                                      n_threads_outer = n_threads_outer,
                                      tile_warm = tile_warm,
                                      prune_tol = prune_tol,
                                      force_sparse = force_sparse,
                                      cell_coupling = cell_coupling,
                                      hessian_pd_mode = hessian_pd_mode,
                                      step_curvature_mode = step_curvature_mode,
                                      inner_refresh = inner_refresh,
                                      x_init_per_cell = x_init_per_cell,
                                      compute_skew = compute_skew,
                                      skew_idx = skew_idx,
                                      fixed_block_p = fixed_block_p,
                                      fixed_block_constraints =
                                          fixed_block_constraints,
                                      debias = debias,
                                      cila = cila,
                                      screen_log_offset = screen_log_offset,
                                      screen_only = screen_only)
    }
}

.joint_backends <- list(
    bym2 = list(
        build_grids = function(prior, has_copy, alpha_axis, phi_axes = NULL) {
            sigma_axis <- prior$sigma_grid %||% .nl_grid_axis("field_sd")
            rho_axis <- prior$rho_grid %||% .nl_grid_axis("bym2_rho")
            .joint_cartesian(list(sigma = sigma_axis, rho = rho_axis),
                              has_copy, alpha_axis, phi_axes)
        },
        call_kernel = .make_joint_kernel_caller("bym2"),
        theta_grid = function(grids, has_copy) {
            base <- if (has_copy) {
                cbind(sigma = grids$sigma, rho = grids$rho,
                      alpha = grids$alpha)
            } else {
                cbind(sigma = grids$sigma, rho = grids$rho)
            }
            .append_phi_columns(base, grids)
        },
        layout = function(arms, prior) {
            .joint_layout(arms, prior$n_spatial_units, n_spatial_blocks = 2L,
                          spatial_block_names = c("phi_start", "theta_start"))
        }
    ),

    icar = list(
        build_grids = function(prior, has_copy, alpha_axis, phi_axes = NULL) {
            sigma_axis <- prior$sigma_grid %||% .nl_grid_axis("field_sd")
            .joint_cartesian(list(sigma = sigma_axis), has_copy,
                              alpha_axis, phi_axes)
        },
        call_kernel = .make_joint_kernel_caller("icar"),
        theta_grid = function(grids, has_copy) {
            base <- if (has_copy) {
                cbind(sigma = grids$sigma, alpha = grids$alpha)
            } else {
                cbind(sigma = grids$sigma)
            }
            .append_phi_columns(base, grids)
        },
        layout = function(arms, prior) {
            .joint_layout(arms, prior$n_spatial_units, n_spatial_blocks = 1L,
                          spatial_block_names = "phi_start")
        }
    ),

    car_proper = list(
        build_grids = function(prior, has_copy, alpha_axis, phi_axes = NULL) {
            sigma_axis   <- prior$sigma_grid %||% .nl_grid_axis("field_sd")
            rho_car_axis <- prior$rho_car_grid %||% .nl_grid_axis("joint_car_rho")
            .joint_cartesian(list(sigma = sigma_axis, rho_car = rho_car_axis),
                              has_copy, alpha_axis, phi_axes)
        },
        call_kernel = .make_joint_kernel_caller("car_proper"),
        theta_grid = function(grids, has_copy) {
            base <- if (has_copy) {
                cbind(sigma = grids$sigma, rho_car = grids$rho_car,
                      alpha = grids$alpha)
            } else {
                cbind(sigma = grids$sigma, rho_car = grids$rho_car)
            }
            .append_phi_columns(base, grids)
        },
        layout = function(arms, prior) {
            .joint_layout(arms, prior$n_spatial_units, n_spatial_blocks = 1L,
                          spatial_block_names = "phi_start")
        }
    )
)

# Single-block call_kernel: route through cpp_nested_laplace_joint_multi by
# packing the legacy prior + per-arm spatial_idx into a length-1
# blocks_spec list and a theta_grid matrix matching the C++ side's axis
# conventions:
#
#   bym2  +copy: axes = (sigma_occ, sigma_pos, rho)   -- unit-precision
#                latent. R side feeds sigma_occ = sigma, sigma_pos =
#                alpha * sigma from the (sigma, alpha) outer grid; the
#                C++ reads two named columns and uses them as per-arm
#                scaling factors. Mathematically identical to the legacy
#                (sigma_occ, sigma_pos) layout cell-by-cell, but the SET
#                of cells covers alpha evenly instead of forming a
#                Cartesian product in sigma_pos.
#   bym2  -copy: axes = (sigma, rho)                  -- sigma rolled into
#                d_fac directly (no copy block).
#   icar  +copy: axes = (sigma_occ, sigma_pos)        -- as above.
#   icar  -copy: axes = (tau,)                        -- tau on prior;
#                grid$sigma is translated to tau = 1 / sigma^2.
#   car_proper +copy: (sigma_occ, sigma_pos, rho_car) -- as above.
#   car_proper -copy: (sigma, rho_car)                -- sigma in d_fac.
#
# The legacy backends used `grid$sigma` for both the donor amplitude
# (copy case) and the prior scale (no-copy case). The C++ kernel is
# unchanged after the reparameterization; only the physical cells it
# visits change.
.joint_call_kernel_via_multi <- function(type, arms, prior, cp, grids,
                                          max_iter, tol, n_threads,
                                          x_init, store_Q, arm_names,
                                          n_threads_outer = 1L,
                                          tile_warm = TRUE,
                                          prune_tol = 0.0,
                                          force_sparse = FALSE,
                                          cell_coupling = "separable",
                                          hessian_pd_mode = 0L,
                                          step_curvature_mode = 0L,
                                          inner_refresh = 1L,
                                          x_init_per_cell = NULL,
                                          compute_skew = FALSE,
                                          skew_idx = NULL,
                                          fixed_block_p = 0L,
                                          fixed_block_constraints = NULL,
                                          debias = NULL,
                                          cila = NULL,
                                          screen_log_offset = NULL,
                                          screen_only = FALSE) {
    n_arms <- length(arms)
    blk <- prior; blk$type <- type
    .nl_check_block_fields(blk, "joint_single")
    missing_spi <- which(vapply(arms, function(a) length(a$spatial_idx) == 0L,
                                logical(1)))
    if (length(missing_spi)) {
        stop("arm(s) ", paste(missing_spi, collapse = ", "),
             " carry no `spatial_idx`; a '", type,
             "' field needs one per arm.", call. = FALSE)
    }
    spi <- lapply(arms, function(a) as.integer(a$spatial_idx))

    block_spec <- list(
        type            = type,
        n_spatial_units = as.integer(prior$n_spatial_units),
        adj_row_ptr     = as.integer(prior$adj_row_ptr),
        adj_col_idx     = as.integer(prior$adj_col_idx),
        n_neighbors     = as.integer(prior$n_neighbors),
        spatial_idx     = spi
    )
    if (type == "bym2") {
        block_spec$scale_factor <- as.numeric(prior$scale_factor %||% 1.0)
        if (!is.null(prior$node_prec)) {
            block_spec$node_prec <- as.numeric(prior$node_prec)
        }
    }

    # Construct theta_grid columns in the order the C++ kernel expects.
    # Names get a `b1.` prefix to match cpp_nested_laplace_joint_multi's
    # naming convention. The C++ axis names `sigma_occ` / `sigma_pos` are
    # preserved as a contract with the kernel; the R outer grid lives in
    # (sigma, alpha) space and materializes sigma_pos = alpha * sigma at
    # the kernel-call boundary.
    if (cp$has_copy) {
        sigma_occ_col <- as.numeric(grids$sigma)
        sigma_pos_col <- as.numeric(grids$sigma) * as.numeric(grids$alpha)
        cols <- switch(
            type,
            bym2       = list(b1.sigma_occ = sigma_occ_col,
                              b1.sigma_pos = sigma_pos_col,
                              b1.rho       = grids$rho),
            icar       = list(b1.sigma_occ = sigma_occ_col,
                              b1.sigma_pos = sigma_pos_col),
            car_proper = list(b1.sigma_occ = sigma_occ_col,
                              b1.sigma_pos = sigma_pos_col,
                              b1.rho_car   = grids$rho_car)
        )
        copy_block <- 0L
    } else {
        cols <- switch(
            type,
            bym2       = list(b1.sigma = grids$sigma, b1.rho = grids$rho),
            icar       = list(b1.tau   = 1.0 / (as.numeric(grids$sigma)^2)),
            car_proper = list(b1.sigma = grids$sigma, b1.rho_car = grids$rho_car)
        )
        copy_block <- -1L
    }
    theta_grid <- do.call(cbind, lapply(cols, as.numeric))
    colnames(theta_grid) <- names(cols)
    axis_offsets <- as.integer(c(0L, length(cols)))

    phi_grid_per_arm <- .joint_phi_grid_per_arm(grids, arm_names)

    # Tile partition for the three-tier warm-start. Tile = unique
    # (sigma, [rho/rho_car,] phi_<arm>...) per cell,
    # i.e. every axis except the copy coefficient alpha. The joint mode is
    # tile-constant on the donor arm (Q + design depend on (sigma, rho)),
    # so any non-alpha cell in the same tile gives a much better warm-
    # start than the global pilot for the inner Newton at boundary alpha
    # cells. Only active when n_threads_outer > 1 and the user opts in
    # (tile_warm = TRUE, the default).
    tile_partition <- NULL
    if (isTRUE(tile_warm) && cp$has_copy &&
        as.integer(n_threads_outer) > 1L &&
        !is.null(grids$alpha) && length(grids$alpha) > 0L) {
        sigma_vec <- as.numeric(grids$sigma)
        n_grid <- length(sigma_vec)
        non_alpha_cols <- list(sigma = sigma_vec)
        if (!is.null(grids$rho)) {
            non_alpha_cols$rho <- as.numeric(grids$rho)
        }
        if (!is.null(grids$rho_car)) {
            non_alpha_cols$rho_car <- as.numeric(grids$rho_car)
        }
        phi_col_names <- grep("^phi_", names(grids), value = TRUE)
        for (nm in phi_col_names) {
            non_alpha_cols[[nm]] <- as.numeric(grids[[nm]])
        }
        non_alpha_mat <- do.call(cbind, non_alpha_cols)
        tile_partition <- .joint_compute_tile_partition(
            non_alpha_mat, as.numeric(grids$alpha), n_grid)
    }

    res <- .cpp_joint_multi(
        arms_list    = arms,
        copy_arm     = as.integer(cp$copy_arm_zero),
        copy_block   = copy_block,
        blocks_spec  = list(block_spec),
        theta_grid   = theta_grid,
        axis_offsets = axis_offsets,
        max_iter     = as.integer(max_iter),
        tol          = as.numeric(tol),
        n_threads    = as.integer(n_threads),
        x_init_nullable = x_init,
        store_Q      = isTRUE(store_Q),
        phi_grid_per_arm = phi_grid_per_arm,
        n_threads_outer = as.integer(n_threads_outer),
        tile_ids        = tile_partition$tile_ids,
        tile_pilot_cells = tile_partition$tile_pilot_cells,
        prune_tol       = as.numeric(prune_tol),
        force_sparse    = isTRUE(force_sparse),
        cell_coupling_name = as.character(cell_coupling),
        hessian_pd_mode = as.integer(hessian_pd_mode),
        step_curvature_mode = as.integer(step_curvature_mode),
        inner_refresh = as.integer(inner_refresh),
        x_init_per_cell = x_init_per_cell,
        compute_skew = isTRUE(compute_skew),
        skew_idx     = skew_idx,
        fixed_block_p = as.integer(fixed_block_p),
        fixed_block_constraints = fixed_block_constraints,
        debias = debias,
        cila = cila,
        screen_log_offset = screen_log_offset,
        screen_only = isTRUE(screen_only)
    )
    # Strip the C++-side theta_grid / axis_offsets -- the backend's
    # `theta_grid()` callback rebuilds them with the user-facing bare
    # names in tulpa_nested_laplace_joint().
    res$theta_grid   <- NULL
    res$axis_offsets <- NULL
    res
}

# --- cheap-pass prune safety gate --------------------------------------------
#
# After a pruned kernel run, decide whether the kept set can be trusted. The
# screen ranked every cell; the full inner Newton ran on the kept cells only, so
# on those the screen's error is measured and on the dropped ones it is not.
# What the gate bounds is the posterior MASS the prune may have dropped
# (`.nl_prune_dropped_mass()`), taking a dropped cell's error to be no worse
# than the worst one measured on a kept cell. Past `.nl_screen("gate_mass")`
# it warns and falls back to the full grid.
#
# The mass, not the ranking, is the question. On a grid placed around its mode
# the neighbouring cells sit within a few nats of each other, so which of them
# the screen puts first says nothing about whether a cell it dropped carried
# weight; and a screen that is off by the same amount on every cell ranks them
# exactly as the full pass does, however large that amount is
# (gcol33/tulpa#925).
#
# `res` is the pruned kernel result (already carries the prune_* fields).
# `resolve_full` is a zero-argument thunk re-running the SAME kernel call with
# `prune_tol = 0` (the full grid). `fn` is the front door the fit came in
# through, so the warning names the function whose `control` the reader set.
# Returns either `res` (gate passed) or the full-grid result (gate tripped),
# each carrying the bound in `prune_dropped_mass_bound`; a tripped gate adds
# `prune_fallback_triggered = TRUE` and `prune_fallback_reason`.
.joint_prune_safety_gate <- function(res, resolve_full,
                                      gate_mass = .nl_screen("gate_mass"),
                                      fn = "tulpa_nested_laplace_joint()") {
    # No prune ran (prune off, prune_tol = 0, or single-cell grid): nothing
    # to gate. The kernel only emits prune_mask when it actually screened.
    if (is.null(res$prune_mask)) return(res)
    # A screened read (a placement pilot's detecting grid) has no full pass to
    # compare the screen against, and is never integrated.
    if (isTRUE(res$prune_screen_only)) return(res)

    bound <- .nl_prune_dropped_mass(res)
    if (is.finite(bound) && bound <= gate_mass) {
        res$prune_dropped_mass_bound <- bound
        return(res)
    }
    reason <- if (is.na(bound)) {
        "the screened result carries no per-cell screen to bound the dropped mass"
    } else if (is.infinite(bound)) {
        "the screen kept no cell with a finite log-marginal"
    } else sprintf(paste0(
        "the cells the screen dropped may carry up to %s of the posterior ",
        "mass, taking their screening error to be the worst one measured on ",
        "the cells it kept"), format(signif(bound, 3)))
    warning(sprintf(paste0(
        "%s: cheap-pass prune is unreliable for this fit (%s); falling back to ",
        "the full grid. Set control$prune = FALSE to skip the screen, or leave ",
        "it -- the full grid is correct."), fn, reason), call. = FALSE)

    full <- resolve_full()
    full$prune_fallback_triggered <- TRUE
    full$prune_fallback_reason    <- reason
    full$prune_dropped_mass_bound <- bound
    full
}

# The share of the posterior a screened grid's dropped cells may carry. With c
# a cell's screened and f its full log-posterior -- the kernel's log-marginal
# plus the hyperprior and cell measure the screen ranked with -- and
# e_max = max over kept cells the screen saw of (f - c), the worst error
# measured,
#
#   B = sum_{j dropped} exp(c_j + e_max) / (R + sum_{k kept} exp(f_k)).
#
# A constant error cancels (B is then the screened share of the dropped cells,
# which the tolerance already bounds); an error that varies across the kept
# cells lifts every dropped cell by the worst of it. A kept cell no screen saw
# (a single-cell refinement call) is mass in the denominator with no error to
# read. R is the mass of the grid a refinement call's cells joined
# (`prune_screen_log_ref`), 0 for a call that is the whole grid. `Inf` when
# nothing was kept, `NA` when the result carries no screen to read.
.nl_prune_dropped_mass <- function(res) {
    mask  <- as.logical(res$prune_mask)
    lm    <- as.numeric(res$log_marginal)
    cheap <- as.numeric(res$prune_cheap_log_marginal)
    n <- length(lm)
    if (length(mask) != n || length(cheap) != n) return(NA_real_)
    mask[is.na(mask)] <- FALSE
    folded <- .nl_log_hyperprior_folded(res, n)
    if (!is.null(folded)) lm <- lm - folded
    off <- res$prune_screen_log_offset
    if (length(off) == n) {
        lm    <- lm + off
        cheap <- cheap + off
    }
    kept <- !mask & is.finite(lm)
    seen <- kept & is.finite(cheap)
    drop <- mask & is.finite(cheap)
    ref  <- as.numeric(res$prune_screen_log_ref %||% -Inf)
    if (!any(seen)) return(Inf)
    if (!any(drop)) return(0)
    e_max <- max(lm[seen] - cheap[seen])
    top   <- max(c(lm[kept], ref))
    sum(exp(cheap[drop] + e_max - top)) /
        (exp(ref - top) + sum(exp(lm[kept] - top)))
}