Nothing
R/nested_laplace_joint_multi.R
In tulpa: Templated Unified Library for Posterior Approximation in Bayesian Hierarchical Models
Defines functions
.joint_posterior_moments_multi
.joint_multi_layout
.joint_fixed_layout
.joint_dispatch_multi
.nlj_multi_inner_skew_at_theta
.joint_attach_pareto_k_multi
.joint_multi_make_kernel_fn
.joint_multi_call_factory
.joint_multi_declared_axes
.joint_multi_hyperprior
.joint_multi_add_hp
.joint_multi_hp_specs
.joint_multi_attach_integration
.joint_multi_hp_families
.joint_multi_measure_specs
.joint_check_multi_hyperpriors
.joint_multi_hp_cols
.joint_block_axis_cols
.joint_hp_named_blocks
.joint_hp_fn_for_block
.joint_parse_hyperprior
.joint_is_block_hyperprior
.joint_hp_field
.joint_multi_phi_per_arm
.joint_multi_cross_phi
.joint_multi_cpp_grid
.multi_block_per_arm_idx
.multi_block_svc_weight
.joint_block_spec_for_cpp
.joint_block_axis_grid
.resolve_copy_multi
.resolve_one_copy_spec
.is_copy_spec_list
.normalise_joint_arm_multi
# Multi-block joint dispatch (list-of-blocks prior).
#
# Part of the joint nested-Laplace driver; the public entry point
# tulpa_nested_laplace_joint() lives in nested_laplace_joint.R.
# =============================================================================
# Multi-block joint dispatch (Phase J-B)
# =============================================================================
#
# Activated when `prior` is a list-of-blocks (each element has a `type`
# field). Routes to cpp_nested_laplace_joint_multi, which builds a
# std::vector<LatentBlock> on the joint side. Each block can be a copy
# block; copy (alpha-coupling) is supported on areal spatial (icar / bym2 /
# car_proper), temporal (rw1 / rw2 / ar1), and unstructured (iid) blocks.
# Validate one arm spec for the multi-block path. Same core as
# `.normalise_joint_arm` *except* `spatial_idx` is no longer required at
# the arm level -- per-arm idx vectors live inside each block spec instead,
# and a missing spatial_idx gets a length-N placeholder of zeros so the
# existing JointArm packaging path doesn't complain.
.normalise_joint_arm_multi <- function(a, k) {
.normalise_joint_arm_core(a, k, spatial_idx = "optional")
}
# Detect whether `copy` is a LIST of copy specs (one coupled field each)
# rather than a single spec. A list-of-specs is an unnamed list whose
# elements are themselves lists carrying a `block` field; a single spec is
# itself a list with a `block` (or `arm`) entry at the top level.
.is_copy_spec_list <- function(copy) {
if (!is.list(copy)) return(FALSE)
if (!is.null(copy$block) || !is.null(copy$arm)) return(FALSE)
nm <- names(copy)
has_no_names <- is.null(nm) || all(!nzchar(nm))
has_no_names && length(copy) > 0L &&
all(vapply(copy, function(s) is.list(s) && !is.null(s$block),
logical(1)))
}
# Resolve and validate one copy spec into a (arm_zero, block_zero,
# alpha_grid) triple. The per-spec validation (arm in range, block in
# range, block type spatial, alpha_grid non-negative non-empty) lives here
# so the single-spec and list-of-specs paths share one source of truth.
.resolve_one_copy_spec <- function(spec, responses, prior_list) {
arm_id <- spec$arm
if (is.null(arm_id)) {
stop("`copy$arm` must be a name or 1-based index.", call. = FALSE)
}
if (is.character(arm_id)) {
nm <- names(responses)
if (is.null(nm) || !arm_id %in% nm) {
stop("`copy$arm` = '", arm_id, "' not found in names(responses).",
call. = FALSE)
}
arm_zero <- match(arm_id, nm) - 1L
} else {
arm_zero <- as.integer(arm_id) - 1L
if (arm_zero < 0L || arm_zero >= length(responses)) {
stop("`copy$arm` index out of range.", call. = FALSE)
}
}
block_id <- spec$block
if (is.null(block_id)) {
stop("`copy$block` must be a 1-based index into `prior` for the ",
"multi-block joint driver.", call. = FALSE)
}
block_zero <- as.integer(block_id) - 1L
if (block_zero < 0L || block_zero >= length(prior_list)) {
stop("`copy$block` index (", block_id, ") out of range for length(prior) (",
length(prior_list), ").", call. = FALSE)
}
block_type <- tolower(prior_list[[block_zero + 1L]]$type)
# Copy (alpha-coupling) is supported on any block whose amplitude is a
# single scalar the donor / copy arms scale independently: areal spatial
# (icar / bym2 / car_proper), temporal (rw1 / rw2 / ar1), unstructured
# (iid), and the separable correlated areal field (mcar). For mcar the copy
# scales the WHOLE correlated (intercept, slope) field by one alpha onto the
# copy arm (the cross-arm transfer); the free cross-covariance Sigma stays
# the within-arm covariance among the fields, integrated over the outer grid
#. Blocks whose per-arm scaling lives elsewhere (lf via
# lambda, hsgp_mo via Sigma) or whose precision is precomputed (tgmrf) do
# not take a copy.
copy_types <- c("icar", "bym2", "car_proper", "rw1", "rw2", "ar1", "iid",
"mcar", "miid")
if (!block_type %in% copy_types) {
stop("`copy$block` points at type '", block_type, "'. Copy semantics ",
"are supported on types (",
paste(shQuote(copy_types), collapse = ", "), "). ",
"Blocks with their own per-arm scaling (lf, hsgp_mo) or a ",
"precomputed precision (tgmrf) do not take a copy.", call. = FALSE)
}
# `alpha_grid` states the nodes; `alpha_n` raises the resolution of the
# engine's own axis, which carries 0 (so the "no copy" base model keeps
# posterior mass) plus a log-spaced slab.
alpha_axis <- .nl_copy_alpha_axis(spec[["alpha_grid"]], spec[["alpha_n"]],
what = "`copy`")
if (length(alpha_axis) == 0L) {
stop("`copy$alpha_grid` must have at least one non-negative value.",
call. = FALSE)
}
if (any(alpha_axis < 0)) {
stop("`copy$alpha_grid` values must be non-negative.",
call. = FALSE)
}
list(arm_zero = arm_zero, block_zero = block_zero, alpha_grid = alpha_axis)
}
# Resolve the copy specification for the multi-block joint driver. `copy`
# may be a single spec `list(arm, block, alpha_grid)` (back-compat) OR an
# unnamed list of such specs (N coupled fields, one alpha axis each). Each
# spec must name a distinct copy-supported block. Returns vectors parallel
# over the copy blocks:
# has_copy, copy_arms_zero (int vec), copy_blocks_zero (int vec),
# alpha_grids (list of numeric, one per copy block).
.resolve_copy_multi <- function(copy, responses, prior_list) {
if (is.null(copy)) {
return(list(has_copy = FALSE, copy_arms_zero = integer(0),
copy_blocks_zero = integer(0), alpha_grids = list()))
}
specs <- if (.is_copy_spec_list(copy)) copy else list(copy)
resolved <- lapply(specs, .resolve_one_copy_spec,
responses = responses, prior_list = prior_list)
copy_blocks_zero <- vapply(resolved, function(r) r$block_zero, integer(1))
copy_arms_zero <- vapply(resolved, function(r) r$arm_zero, integer(1))
alpha_grids <- lapply(resolved, function(r) r$alpha_grid)
if (anyDuplicated(copy_blocks_zero)) {
stop("Two copy specs target the same block (1-based index ",
paste(unique(copy_blocks_zero[duplicated(copy_blocks_zero)]) + 1L,
collapse = ", "),
"). Each coupled field must name a distinct block.",
call. = FALSE)
}
list(has_copy = TRUE, copy_arms_zero = copy_arms_zero,
copy_blocks_zero = copy_blocks_zero, alpha_grids = alpha_grids)
}
# Per-block axis grid for the multi-block joint driver. When the block is a
# copy block (any of icar / bym2 / car_proper / rw1 / rw2 / ar1 / iid), the
# parameterisation uses (sigma, alpha[, rho/rho_car]) directly. Non-copy
# blocks use the standard single-arm conventions and reuse the `.NL_REGISTRY`
# defaults.
#
# Returns:
# $grid : matrix [n_block_cells x n_axes_for_block]
# $names : axis names
# $prepared: block spec with defaults filled in (so downstream code can
# read prior$adj_row_ptr etc. without re-checking presence)
.joint_block_axis_grid <- function(p, is_copy, alpha_grid,
block_index) {
type <- tolower(p$type)
if (is_copy && type %in% c("mcar", "miid")) {
# Copied correlated field (mcar: areal; miid: per-group RE) -- the
# block's own axes are the p(p+1)/2 log-Cholesky coordinates of Sigma
# (the within-arm covariance among the fields); the copy appends ONE
# trailing `alpha`
# axis carrying the cross-arm copy amplitude. The kernel reads the
# log-Cholesky axes at axis0 and alpha at axis0 + p(p+1)/2 (the last
# column), so alpha MUST be appended last. Sigma stays a free outer-grid
# quantity (NOT a copy-materialized sigma_occ/sigma_pos pair); alpha is
# applied directly via arm_scale, not multiplied by a sigma.
base <- .nl_block_axis_grid(p)
bg <- base$grid
ag <- as.numeric(alpha_grid)
idx <- expand.grid(b = seq_len(nrow(bg)), a = seq_along(ag),
KEEP.OUT.ATTRS = FALSE)
grid <- cbind(bg[idx$b, , drop = FALSE], alpha = ag[idx$a])
colnames(grid) <- c(colnames(bg), "alpha")
return(list(grid = as.matrix(grid), names = colnames(grid),
prepared = base$prepared))
}
if (is_copy) {
sigma_axis <- as.numeric(p$sigma_grid %||% .nl_grid_axis("field_sd"))
# Copy axes always lead with (sigma, alpha); the kernel materializes
# sigma_donor = sigma and sigma_copy = alpha * sigma in those two
# leading columns (see .joint_multi_cpp_grid), so any trailing
# correlation axis (rho / rho_car) must follow them. A unit-precision
# block (icar / rw1 / rw2 / iid) has just (sigma, alpha); a block with a
# second hyperparameter (bym2 rho, car_proper rho_car, ar1 rho) appends
# it.
if (type %in% c("icar", "rw1", "rw2", "iid")) {
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else if (type == "bym2") {
rho <- p$rho_grid %||% .nl_grid_axis("bym2_rho")
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
rho = as.numeric(rho),
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else if (type == "car_proper") {
rho_car <- p$rho_car_grid %||% .nl_grid_axis("joint_car_rho")
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
rho_car = as.numeric(rho_car),
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else if (type == "ar1") {
rho <- p$rho_grid %||% .nl_grid_axis("ar1_rho")
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
rho = as.numeric(rho),
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else {
stop("Block ", block_index, ": copy semantics are not supported ",
"for type '", type, "'.", call. = FALSE)
}
return(list(grid = as.matrix(gr),
names = colnames(gr),
prepared = p))
}
# Non-copy: reuse the single-arm registry's grid construction. The
# registry returns the standard (tau, [rho]) / (sigma, [rho]) axes.
.nl_block_axis_grid(p)
}
# Convert one R-side block to the C++ `cpp_nested_laplace_joint_multi`
# format. Per-arm idx vectors live in the BLOCK spec (not the arm spec);
# the user supplies `spatial_idx` / `temporal_idx` / `obs_idx` as a list
# of length n_arms.
#
# `arms` is the list of normalised arm specs produced by
# `.normalise_joint_arm_multi()`. It is consumed by block types that need
# to read covariate columns at fit-time (HSGP with `svc_column`). Pass
# NULL only from contexts where no such block can appear.
.joint_block_spec_for_cpp <- function(p, n_arms, block_index, arms = NULL) {
type <- tolower(p$type)
projected <- .nl_block_is_projected(p)
.nl_check_block_fields(
p, c("axis", if (projected) "projected" else "joint"), block_index)
arm_n_obs <- if (is.null(arms)) NULL
else vapply(arms, function(a) length(a$y), integer(1))
if (type %in% .NL_PROJECTABLE_AREAL) {
spatial_idx <- if (!projected)
.multi_block_per_arm_idx(p$spatial_idx, n_arms, block_index,
"spatial_idx", arm_n_obs,
p$n_spatial_units)
out <- list(
type = type,
spatial_idx = spatial_idx,
n_spatial_units = as.integer(p$n_spatial_units),
adj_row_ptr = as.integer(p$adj_row_ptr),
adj_col_idx = as.integer(p$adj_col_idx),
n_neighbors = as.integer(p$n_neighbors)
)
if (projected) {
out$projector <- .nl_block_projector(p, n_arms, block_index,
arm_n_obs)
}
if (type == "bym2") {
out$scale_factor <- as.numeric(p$scale_factor %||% 1.0)
if (!is.null(p$node_prec)) out$node_prec <- as.numeric(p$node_prec)
}
# Optional per-arm per-row design weight (areal SVC). When set, the
# field's contribution to arm k row i is row-scaled by svc_weight[[k]][i]:
# eta_i += svc_weight[[k]][i] * amplitude * z[cell_i].
# Parallel to `spatial_idx`: a list of length n_arms, each a numeric
# vector matching that arm's obs count and the per-arm spatial_idx
# length. Composes multiplicatively with copy `arm_scale`.
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, spatial_idx, n_arms, block_index)
# Aliased fixed-effect column of a weighted (areal SVC) field: its
# constant folds into this coefficient (0 = intercept, or the
# covariate's column) rather than the intercept; -1 = no counterpart
# (the level is left in the field). Absent => no fold.
if (!is.null(p$svc_beta_offset)) {
out$svc_beta_offset <- as.integer(p$svc_beta_offset)
}
}
# Replicated CAR (`by =`): number of disjoint equal-size components in the
# block-diagonal graph (the intrinsic ICAR kernel pins each per component
# and uses (n - L) in its rank-deficiency normalizer). Absent = connected.
if (!is.null(p$n_components)) {
out$n_components <- as.integer(p$n_components)
}
out
} else if (type == "mcar") {
# Separable multivariate CAR: p coupled areal fields sharing Sigma (x)
# Q^-1. One block over p * n_spatial_units latent. spatial_idx is the
# per-arm 1-based cell index (shared by all fields); field_weight is the
# per-field design column (the intercept's is all-ones, a covariate
# column is the per-row value) -- outer length p, inner length n_arms.
spatial_idx <- .multi_block_per_arm_idx(p$spatial_idx, n_arms,
block_index, "spatial_idx",
arm_n_obs, p$n_spatial_units)
n_fields <- as.integer(p$n_fields)
if (length(p$field_weight) != n_fields) {
stop("Block ", block_index, " (type 'mcar'): `field_weight` must be ",
"a list of length n_fields (", n_fields, ").", call. = FALSE)
}
field_weight <- lapply(seq_len(n_fields), function(a) {
.multi_block_svc_weight(p$field_weight[[a]], spatial_idx, n_arms,
block_index)
})
out <- list(
type = "mcar",
spatial_idx = spatial_idx,
n_spatial_units = as.integer(p$n_spatial_units),
n_fields = n_fields,
adj_row_ptr = as.integer(p$adj_row_ptr),
adj_col_idx = as.integer(p$adj_col_idx),
n_neighbors = as.integer(p$n_neighbors),
field_weight = field_weight
)
# Per-field aliased fixed-effect column (length n_fields): field a's
# global constant folds into field_beta_offset[a] (0 = intercept for the
# all-ones field, the covariate's column otherwise, -1 = no counterpart).
# Absent => no folds (the augmentation identifies every level in place).
if (!is.null(p$field_beta_offset)) {
if (length(p$field_beta_offset) != n_fields) {
stop("Block ", block_index, " (type 'mcar'): `field_beta_offset` ",
"must have length n_fields (", n_fields, ").", call. = FALSE)
}
out$field_beta_offset <- as.integer(p$field_beta_offset)
}
# Replicated MCAR (`by =`): L equal-size components per field.
if (!is.null(p$n_components)) {
out$n_components <- as.integer(p$n_components)
}
out
} else if (type %in% c("rw1", "rw2", "ar1")) {
temporal_idx <- .multi_block_per_arm_idx(p$temporal_idx, n_arms,
block_index, "temporal_idx",
arm_n_obs, p$n_times)
out <- list(
type = type,
temporal_idx = temporal_idx,
n_times = as.integer(p$n_times)
)
if (type %in% c("rw1", "rw2")) out$cyclic <- isTRUE(p$cyclic)
# Optional per-arm per-row design weight (temporal SVC). Parallel to the
# areal branch: the field's contribution to arm k row i is row-scaled by
# svc_weight[[k]][i] (eta_i += svc_weight[[k]][i] * amplitude *
# f[time_i]), used for a temporally varying slope on a covariate column.
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, temporal_idx, n_arms, block_index)
# Aliased fixed-effect column of a weighted (temporal SVC / TVC)
# field, as in the areal branch above; absent => no fold.
if (!is.null(p$svc_beta_offset)) {
out$svc_beta_offset <- as.integer(p$svc_beta_offset)
}
}
out
} else if (type == "iid") {
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, p$n_units)
out <- list(
type = "iid",
obs_idx = obs_idx,
n_units = as.integer(p$n_units)
)
# Optional per-arm per-row design weight (random slope on a covariate
# column). When set, the field's contribution to arm k row i is
# row-scaled by svc_weight[[k]][i]:
# eta_i += svc_weight[[k]][i] * sigma * u[obs_idx_i].
# Parallel to the areal / temporal SVC branches: a list of length n_arms,
# each a numeric vector matching that arm's obs count (the per-arm obs_idx
# length). Composes multiplicatively with copy `arm_scale`. An
# uncorrelated slope (x || g) / (0 + x | g) is one weighted iid block per
# coefficient, each with its own sigma axis.
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, obs_idx, n_arms, block_index)
}
out
} else if (type == "miid") {
# Multivariate IID: p coupled per-group coefficient
# fields sharing a free Sigma (the Q = I sibling of mcar). One block over
# p * n_groups latent. `obs_idx` is the per-arm 1-based grouping index
# (the mcar `spatial_idx` role); `field_weight` is the per-coefficient
# design column (intercept all-ones, slope = covariate) -- outer length
# n_fields, inner length n_arms. Expresses a correlated random slope
# (1 + x | g): n_fields = 1 + n_slopes.
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, p$n_groups)
n_fields <- as.integer(p$n_fields)
if (length(p$field_weight) != n_fields) {
stop("Block ", block_index, " (type 'miid'): `field_weight` must be ",
"a list of length n_fields (", n_fields, ").", call. = FALSE)
}
field_weight <- lapply(seq_len(n_fields), function(a) {
.multi_block_svc_weight(p$field_weight[[a]], obs_idx, n_arms,
block_index)
})
list(
type = "miid",
obs_idx = obs_idx,
n_groups = as.integer(p$n_groups),
n_fields = n_fields,
field_weight = field_weight
)
} else if (type == "tgmrf") {
# User-supplied GMRF: per-grid CSC Q + logdet + log p(theta).
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, p$n_latent)
list(
type = "tgmrf",
n_latent = as.integer(p$n_latent),
obs_idx = obs_idx,
Q_csc_p_per_grid = lapply(p$Q_csc_p_per_grid, as.integer),
Q_csc_i_per_grid = lapply(p$Q_csc_i_per_grid, as.integer),
Q_csc_x_per_grid = lapply(p$Q_csc_x_per_grid, as.numeric),
logdet_Q_per_grid = as.numeric(p$logdet_Q_per_grid),
log_prior_theta_per_grid = as.numeric(p$log_prior_theta_per_grid)
)
} else if (type == "nngp") {
spatial_idx <- if (!projected)
.multi_block_per_arm_idx(p$spatial_idx, n_arms, block_index,
"spatial_idx", arm_n_obs, p$n_spatial)
out <- c(list(type = "nngp", spatial_idx = spatial_idx),
.nl_nngp_field_args(p))
if (projected) {
out$projector <- .nl_block_projector(p, n_arms, block_index,
arm_n_obs)
}
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, spatial_idx, n_arms, block_index)
}
out
} else if (type == "hsgp") {
# HSGP block: shared eigenvalues + per-arm Phi basis matrices.
# Axes are (sigma2, lengthscale).
#
# Optional `svc_column`: 1-based index into each arm's `X`. When
# set, the block becomes a spatially-varying coefficient:
# eta_i += X[i, svc_column] * f(s_i)
# Implemented by row-scaling the per-arm basis at fit-time
# phi_scaled[k][i, m] = phi[k][i, m] * X[k][i, svc_column]
# and passing the scaled phi to the existing DENSE_BASIS factory.
# Multi-scale composition (1.6c) falls out by declaring multiple
# type='hsgp' blocks in the same prior_list.
if (!is.list(p$phi) || length(p$phi) != n_arms) {
stop("Block ", block_index, " (type 'hsgp'): `phi` must be a ",
"list of length n_arms (", n_arms, ").", call. = FALSE)
}
phi <- lapply(p$phi, function(M) {
M <- as.matrix(M); storage.mode(M) <- "double"; M
})
if (!is.null(p$svc_column)) {
if (is.null(arms)) {
stop("Block ", block_index, " (type 'hsgp'): `svc_column` ",
"requires the joint multi-block dispatcher path ",
"(arms not available here).", call. = FALSE)
}
svc <- as.integer(p$svc_column)
if (length(svc) != 1L || is.na(svc) || svc < 1L) {
stop("Block ", block_index, " (type 'hsgp'): `svc_column` ",
"must be a single positive integer (1-based column ",
"index into each arm's X).", call. = FALSE)
}
for (k in seq_len(n_arms)) {
Xk <- arms[[k]]$X
if (svc > ncol(Xk)) {
stop("Block ", block_index, " (type 'hsgp'): ",
"`svc_column` = ", svc, " exceeds ncol(arm ", k,
"$X) = ", ncol(Xk), ".", call. = FALSE)
}
if (nrow(phi[[k]]) != nrow(Xk)) {
stop("Block ", block_index, " (type 'hsgp'): ",
"nrow(phi[[", k, "]]) (", nrow(phi[[k]]),
") must equal nrow(arm ", k, "$X) (", nrow(Xk),
") when svc_column is set.", call. = FALSE)
}
phi[[k]] <- phi[[k]] * as.numeric(Xk[, svc])
}
}
list(
type = "hsgp",
m_total = as.integer(p$m_total),
phi = phi,
n_obs_per_arm = as.integer(p$n_obs_per_arm),
eigenvalues = as.numeric(p$eigenvalues)
)
} else if (type == "hsgp_mo") {
# Multi-output (co-regionalization) HSGP block. K = n_arms
# correlated latent fields share basis + eigenvalues; coefficients are
# cross-output-correlated via Sigma. First ship restricts K == 2 with
# axes (sigma_1, sigma_2, rho, ell), all raw.
if (n_arms != 2L) {
stop("Block ", block_index, " (type 'hsgp_mo'): ",
"requires n_arms == 2 (got ", n_arms, "). Multi-output HSGP ",
"with K > 2 needs an LKJ-cholesky correlation ",
"parameterization on the outer grid, which is not supported.",
call. = FALSE)
}
if (!is.list(p$phi) || length(p$phi) != n_arms) {
stop("Block ", block_index, " (type 'hsgp_mo'): `phi` must be a ",
"list of length n_arms (", n_arms, ").", call. = FALSE)
}
phi <- lapply(p$phi, function(M) {
M <- as.matrix(M); storage.mode(M) <- "double"; M
})
list(
type = "hsgp_mo",
m_total = as.integer(p$m_total),
phi = phi,
n_obs_per_arm = as.integer(p$n_obs_per_arm),
eigenvalues = as.numeric(p$eigenvalues)
)
} else if (type == "spde") {
# SPDE block: per-arm sparse barycentric projector A and shared FEM
# mesh (C0_diag, G1). Axes are (range, sigma). Optional rational
# coefficients for fractional nu.
per_arm <- function(field) {
v <- p[[field]]
if (!is.list(v) || length(v) != n_arms) {
stop("Block ", block_index, " (type 'spde'): `", field,
"` must be a list of length n_arms (", n_arms, ").",
call. = FALSE)
}
v
}
out <- list(
type = "spde",
n_mesh = as.integer(p$n_mesh),
A_x = lapply(per_arm("A_x"), as.numeric),
A_i = lapply(per_arm("A_i"), as.integer),
A_p = lapply(per_arm("A_p"), as.integer),
n_obs_per_arm = as.integer(p$n_obs_per_arm),
C0_diag = as.numeric(p$C0_diag),
G1_x = as.numeric(p$G1_x),
G1_i = as.integer(p$G1_i),
G1_p = as.integer(p$G1_p),
nu = as.numeric(p$nu)
)
if (!is.null(p$rational_poles))
out$rational_poles <- as.numeric(p$rational_poles)
if (!is.null(p$rational_weights))
out$rational_weights <- as.numeric(p$rational_weights)
out
} else if (type == "lf") {
# Latent factor block. Per-arm `obs_idx` maps each
# obs to a slot in the factor field `u` (length n_latent); the
# per-arm loadings `lambda` (length n_arms) are part of the joint
# latent vector and co-optimized by the inner Newton. No outer-
# grid axes -- identifiability handled by tight Gaussian anchors
# on (u_1, lambda_1) inside the C++ factory.
#
# The unit count is checked BEFORE the indices are checked against it:
# an `n_latent` of 1 makes every real index out of range, so validating
# the indices first reports the block's shape error as an index error.
# The C++ factory keeps its own copy of this check for a consumer that
# builds the block through `LinkingTo` rather than this door.
n_latent <- as.integer(p$n_latent)
if (length(n_latent) != 1L || is.na(n_latent) || n_latent < 2L) {
stop("Block ", block_index, " (type 'lf'): `n_latent` must be at ",
"least 2 (got ", format(p$n_latent), "). The factor field ",
"needs a second slot for the anchor on (u_1, lambda_1) to ",
"leave anything free.", call. = FALSE)
}
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, n_latent)
out <- list(
type = "lf",
n_latent = n_latent,
obs_idx = obs_idx
)
if (!is.null(p$sigma_u)) out$sigma_u <- as.numeric(p$sigma_u)
if (!is.null(p$sigma_lambda)) out$sigma_lambda <- as.numeric(p$sigma_lambda)
if (!is.null(p$anchor_eps)) out$anchor_eps <- as.numeric(p$anchor_eps)
out
} else {
stop("Block type '", type, "' is not supported in multi-block joint priors.",
call. = FALSE)
}
}
# Coerce + validate the optional `svc_weight` for an areal block into a
# length-n_arms list of numeric vectors. Each entry must match that arm's
# obs count, which is the length of that arm's resolved `spatial_idx`.
# Accepts either a list (already per-arm) or a single numeric vector
# (replicated across arms when every arm shares the same per-row weight).
.multi_block_svc_weight <- function(svc_weight, spatial_idx, n_arms,
block_index) {
if (is.list(svc_weight)) {
if (length(svc_weight) != n_arms) {
stop("Block ", block_index, ": `svc_weight` is a list of ",
length(svc_weight), " entries but n_arms = ", n_arms, ".",
call. = FALSE)
}
w_list <- lapply(svc_weight, as.numeric)
} else {
w_list <- rep(list(as.numeric(svc_weight)), n_arms)
}
for (k in seq_len(n_arms)) {
n_k <- length(spatial_idx[[k]])
if (length(w_list[[k]]) != n_k) {
stop("Block ", block_index, ": `svc_weight[[", k, "]]` has length ",
length(w_list[[k]]), " but arm ", k, " has ", n_k,
" observations (matching its `spatial_idx`).", call. = FALSE)
}
if (any(!is.finite(w_list[[k]]))) {
stop("Block ", block_index, ": `svc_weight[[", k, "]]` contains ",
"non-finite values.", call. = FALSE)
}
}
w_list
}
# Coerce a per-arm idx field into a length-n_arms list of integer vectors.
# Accepts either a list (already per-arm) or a single vector (replicated
# across arms; useful when every arm shares the same indexing -- e.g. years
# 1..T mapped one-to-one).
# `arm_n_obs`, when supplied, is each arm's observation count
# (length(arms[[k]]$y)): the C++ side (make_per_arm_idx_fn,
# src/nested_laplace_joint_multi.cpp) reads `cache[k_arm][i]` for every row i
# of that arm with NO bounds check, so a per-arm entry shorter than that arm's
# row count is a read out of bounds -- an access violation, not a clean error
# -- the moment a grid cell scatters that arm's rows into this block. There is
# no "this block skips arm k" shorthand via a short/empty vector: a block
# that should not affect an arm still needs a full-length index vector for
# it; exclude the arm's CONTRIBUTION instead (e.g. a field_coef of 0 on that
# arm), which the engine already tests via `d_eff_cache == 0`.
#
# `n_units`, when supplied, is the block's own latent unit count
# (n_spatial_units / n_times / n_units / n_groups / n_latent). Every index must
# land in 1..n_units. The kernel treats an out-of-range index as "this row does
# not see this block" and moves on at every one of its six consumers, so a
# 0-based vector, an off-by-one, or an NA does not fail: the affected rows drop
# out of eta, out of the gradient and out of the Hessian, the inner Newton
# converges on the reduced problem, `converged` is TRUE, and the log-marginal
# surface is the one for a different model. The range is the only place that
# distinguishes a deliberately absent row from a wrong one, so it is checked
# here, once, where the length already is.
.multi_block_per_arm_idx <- function(idx, n_arms, block_index, field_name,
arm_n_obs = NULL, n_units = NULL) {
out <- if (is.list(idx)) {
if (length(idx) != n_arms) {
stop("Block ", block_index, ": `", field_name, "` is a list of ",
length(idx), " entries but n_arms = ", n_arms, ".",
call. = FALSE)
}
lapply(idx, as.integer)
} else {
rep(list(as.integer(idx)), n_arms)
}
if (!is.null(arm_n_obs)) {
for (k in seq_len(n_arms)) {
if (length(out[[k]]) != arm_n_obs[k]) {
stop("Block ", block_index, ": `", field_name, "[[", k, "]]` has ",
length(out[[k]]), " entries but arm ", k, " has ",
arm_n_obs[k], " observations -- every arm needs a ",
"matching-length index vector for every block (there is ",
"no arm-exclusion shorthand). To exclude arm ", k,
" from this block, zero its CONTRIBUTION (e.g. a ",
"`field_coef` of 0 on that arm) rather than shortening ",
"`", field_name, "`.", call. = FALSE)
}
}
}
if (!is.null(n_units)) {
n_units <- as.integer(n_units)
if (is.na(n_units) || n_units < 1L) {
stop("Block ", block_index, ": the latent unit count backing `",
field_name, "` must be a positive integer, got ", n_units, ".",
call. = FALSE)
}
for (k in seq_len(n_arms)) {
v <- out[[k]]
if (length(v) == 0L) next
# The kernel's contract for a block index is that `l > 0 && l <=
# size` contributes and anything else is skipped (LatentBlock::idx,
# and every scatter that reads it). A non-positive entry is
# therefore the documented "this observation does not reach this
# block" sentinel -- what an arm-exclusive block spells on the arms
# it does not reach -- and is accepted. An entry ABOVE n_units names
# a latent unit that does not exist: the kernel drops it just as
# silently, and there is no reading under which that was meant.
bad <- is.na(v) | v > n_units
if (any(bad)) {
i <- which(bad)[1L]
stop("Block ", block_index, ": `", field_name, "[[", k, "]][",
i, "]` is ", v[i], ", above the block's unit count ",
n_units, " (", sum(bad), " of ", length(v),
" entries are). Every positive index must name one of the ",
"block's latent units; an out-of-range one is silently ",
"dropped by the kernel, which fits a different model and ",
"reports convergence. Use 0 (or any non-positive value) ",
"for an observation that does not reach this block.",
call. = FALSE)
}
# A COMPLETE 0-based enumeration is the one pattern likelier to be
# an off-by-one than a sentinel: every unit but the last named, and
# the unnamed slot filled by the value that also means "skip". A
# sentinel pattern reaching exactly 1..n_units-1 and never n_units
# is possible, so this warns rather than rejecting.
if (n_units > 1L && any(v == 0L) &&
setequal(unique(v), seq.int(0L, n_units - 1L))) {
warning("Block ", block_index, ": `", field_name, "[[", k,
"]]` takes every value in 0..", n_units - 1L,
" and never ", n_units, ", the signature of a 0-based ",
"index vector. Indices are 1-based and 0 means \"this ",
"observation does not reach this block\", so unit ",
n_units, " would go unreached and every other ",
"observation would land one unit low.", call. = FALSE)
}
}
}
out
}
# Materialize the C++-facing theta_grid from a user-facing multi-block
# `joint_grid` (the latent-block axis columns; any trailing phi_<arm>
# dispersion columns are excluded here -- they ride on `phi_grid_per_arm`).
# The kernel reads `sigma_occ` / `sigma_pos` on a copy block; the R-side grid
# lives in (sigma, alpha), so this materializes sigma_occ = sigma and
# sigma_pos = alpha * sigma per copy block. Non-copy block columns pass
# through unchanged. Single source of truth for the main dispatch and the
# Pareto-k re-evaluation.
.joint_multi_cpp_grid <- function(joint_grid, axis_offsets, B, cp) {
cpp_grid <- joint_grid[, seq_len(axis_offsets[B + 1L]), drop = FALSE]
if (!cp$has_copy) return(cpp_grid)
new_names <- colnames(cpp_grid)
for (b_copy in (cp$copy_blocks_zero + 1L)) {
cols_b <- (axis_offsets[b_copy] + 1L):axis_offsets[b_copy + 1L]
bare_b <- sub("^b[0-9]+\\.", "", colnames(cpp_grid)[cols_b])
i_sigma <- match("sigma", bare_b)
i_alpha <- match("alpha", bare_b)
# A copied correlated field (mcar) has no scalar
# `sigma` axis -- its amplitude lives in the free Sigma (the log-Cholesky
# axes), and the copy applies alpha directly via arm_scale (NOT
# sigma_occ = sigma / sigma_pos = alpha*sigma). The alpha column passes
# through to the kernel unchanged; the log-Cholesky axes pass through too.
if (is.na(i_sigma) && !is.na(i_alpha)) next
if (is.na(i_sigma) || is.na(i_alpha)) {
stop(".joint_multi_cpp_grid: copy block missing 'sigma' or 'alpha' axis.",
call. = FALSE)
}
sigma_col <- as.numeric(cpp_grid[, cols_b[i_sigma]])
alpha_col <- as.numeric(cpp_grid[, cols_b[i_alpha]])
cpp_grid[, cols_b[i_sigma]] <- sigma_col # -> sigma_occ
cpp_grid[, cols_b[i_alpha]] <- alpha_col * sigma_col # -> sigma_pos
new_names[cols_b[i_sigma]] <- paste0("b", b_copy, ".sigma_occ")
new_names[cols_b[i_alpha]] <- paste0("b", b_copy, ".sigma_pos")
}
colnames(cpp_grid) <- new_names
cpp_grid
}
# Cross the active per-arm dispersion axes of `phi_axes` (the
# `.normalise_phi_grid()` list) onto a latent outer grid as a Cartesian product:
# every latent row repeated once per dispersion cell, latent rows varying
# fastest, the dispersion columns appended as `phi_<arm>`. Returns the crossed
# grid and `latent_row`, the latent row each crossed cell repeats. The cell
# layout depends only on the latent row count and each axis's node count, so
# two responses whose axes share node counts share it, whatever the nodes are.
.joint_multi_cross_phi <- function(joint_grid, phi_axes, axis_names) {
active <- phi_axes[vapply(phi_axes, length, integer(1)) > 0L]
phi_extra <- do.call(expand.grid,
c(active, list(KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)))
joint_idx <- expand.grid(joint = seq_len(nrow(joint_grid)),
phi = seq_len(nrow(phi_extra)),
KEEP.OUT.ATTRS = FALSE)
grid <- cbind(joint_grid[joint_idx$joint, , drop = FALSE],
as.matrix(phi_extra[joint_idx$phi, , drop = FALSE]))
colnames(grid) <- c(axis_names, paste0("phi_", names(active)))
list(grid = grid, latent_row = joint_idx$joint)
}
# Extract the per-arm dispersion override list from a `joint_grid` carrying
# `phi_<arm>` columns. Returns NULL when no such column exists (the kernel
# then uses each arm's parse-time scalar phi); otherwise a length-n_arms list
# with entry k the column `phi_<arm_names[k]>` or NULL.
.joint_multi_phi_per_arm <- function(joint_grid, arm_names) {
cn <- colnames(joint_grid)
if (is.null(cn) || !any(startsWith(cn, "phi_"))) return(NULL)
out <- vector("list", length(arm_names))
for (k in seq_along(arm_names)) {
col <- paste0("phi_", arm_names[k])
if (col %in% cn) out[[k]] <- as.numeric(joint_grid[, col])
}
out
}
# --- regularizing hyperpriors on a multi-block grid ------------------------
#
# `prior_sigma` and `prior_alpha` name a ROLE -- the donor field amplitude and
# the copy coefficient -- and a multi-block grid can carry that role on several
# blocks at once: an occurrence field copied onto two arms carries `b1.alpha`
# and `b2.alpha`. Two shapes state which axes one density reaches.
#
# * ONE spec, `list(<family>, <params>)`: the density applies to EVERY block
# axis of that role. That is the reading `prior_phi` already takes across
# the per-arm dispersion axes, and the one a single-block grid's single
# axis cannot distinguish. Each axis carries its own copy of the density
# and each axis's atom / continuum split is normalised against that axis's
# own cell widths -- `.hyper_atom_fold_scale()` reads one spec's own
# `log_prior` against its own continuum -- so N axes take N independent
# applications and no normalizer is shared between blocks.
# * A LIST of per-block entries, `list(list(block = , prior = ), ...)`, keyed
# by 1-based block index the way a multi-block `copy` spec keys its blocks.
# A block the list does not name carries no density on that axis.
#
# Parsing returns NULL (flat), one density function (the first shape), or a
# `tulpa_block_hyperprior`: density functions named by block index.
# One named element of a list, NULL where the list carries no such name.
.joint_hp_field <- function(x, name) {
i <- match(name, names(x))
if (is.na(i)) NULL else x[[i]]
}
# Is `spec` the per-block shape? Every element of the scalar shape is a family
# name or a parameter vector; every element of this one is a list naming a
# block.
.joint_is_block_hyperprior <- function(spec) {
if (!is.list(spec) || length(spec) == 0L) return(FALSE)
all(vapply(spec, function(e) is.list(e) && "block" %in% names(e),
logical(1)))
}
# Parse a `prior_sigma` / `prior_alpha` argument. `multi_block` says whether
# the fit's `prior` is a list of blocks: the per-block shape names blocks of
# such a list, so on the single-block path it is refused rather than read as a
# malformed family spec.
.joint_parse_hyperprior <- function(spec, axis_label, multi_block = FALSE) {
if (is.null(spec)) return(NULL)
if (!.joint_is_block_hyperprior(spec)) {
return(.joint_parse_sigma_prior(spec, axis_label))
}
if (!isTRUE(multi_block)) {
stop("`", axis_label, "` is a per-block list, which names blocks of a ",
"list-of-blocks `prior`; this fit takes the single-block joint ",
"path, where the axis is unique. Pass the spec itself, e.g. ",
"list(\"pc.prec\", c(U, alpha)).", call. = FALSE)
}
out <- list()
for (i in seq_along(spec)) {
e <- spec[[i]]
b <- .joint_hp_field(e, "block")
if (length(b) != 1L || !is.numeric(b) || !is.finite(b) ||
b != round(b) || b < 1) {
stop("`", axis_label, "[[", i, "]]$block` must be a single 1-based ",
"block index into `prior`.", call. = FALSE)
}
key <- as.character(as.integer(b))
if (key %in% names(out)) {
stop("`", axis_label, "` names block ", key, " twice. Each block ",
"takes at most one hyperprior on a given axis.", call. = FALSE)
}
p <- .joint_hp_field(e, "prior")
if (is.null(p)) {
stop("`", axis_label, "[[", i, "]]` names block ", key, " but ",
"carries no `prior` field. Each entry is ",
"list(block = <index>, prior = list(<family>, <params>)).",
call. = FALSE)
}
out[[key]] <- .joint_parse_sigma_prior(
p, sprintf("%s[[%d]]$prior", axis_label, i))
}
structure(out, class = "tulpa_block_hyperprior")
}
# The density one hyperprior applies to block `b`: the stated density wherever
# the caller stated one, that block's own entry where the caller keyed the
# prior per block, NULL where neither reaches this block.
.joint_hp_fn_for_block <- function(fn, b) {
if (is.null(fn)) return(NULL)
if (is.function(fn)) return(fn)
.joint_hp_field(fn, as.character(b))
}
# The block indices a per-block hyperprior names. A stated single density names
# none -- it reaches every block carrying the axis -- so this is `integer(0)`
# for it as well as for an absent prior.
.joint_hp_named_blocks <- function(fn) {
if (is.null(fn) || is.function(fn)) return(integer(0))
as.integer(names(fn))
}
# The grid column carrying one axis role in each block that has it, named by
# 1-based block index. A block with no axes at all (a latent-factor block
# carries a 1 x 0 grid) contributes none: its column span is empty, which
# `seq.int(length.out =)` reads as empty where a `from:to` range would run
# backwards into the neighbouring block's columns.
.joint_block_axis_cols <- function(axis_names, axis_offsets, B, role) {
out <- integer(0)
for (b_idx in seq_len(B)) {
n_b <- axis_offsets[b_idx + 1L] - axis_offsets[b_idx]
if (n_b <= 0L) next
cols_b <- seq.int(axis_offsets[b_idx] + 1L, length.out = n_b)
i <- match(role, sub("^b[0-9]+\\.", "", axis_names[cols_b]))
if (is.na(i)) next
out[as.character(b_idx)] <- cols_b[i]
}
out
}
# Which columns of a multi-block `joint_grid` each hyperprior applies to, and
# with which density. One entry per applied (column, role): `col` indexes the
# grid, `role` is `"sigma"` / `"alpha"` / `"phi"`, `block` is the 1-based block
# index (NA for a per-arm phi axis) and `fn` the density folded onto it.
# Single source of the mapping, so the fold below and the axis quadrature
# (`.joint_multi_hp_specs()`) name the same columns and the same densities.
.joint_multi_hp_cols <- function(joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi = NULL) {
cn <- colnames(joint_grid)
out <- list()
fns <- list(sigma = fn_sigma, alpha = fn_alpha)
for (role in c("sigma", "alpha")) {
if (is.null(fns[[role]])) next
cols <- .joint_block_axis_cols(cn, axis_offsets, B, role)
for (k in seq_along(cols)) {
b_idx <- as.integer(names(cols)[k])
fn <- .joint_hp_fn_for_block(fns[[role]], b_idx)
if (is.null(fn)) next
out[[length(out) + 1L]] <- list(col = as.integer(cols[[k]]),
role = role, block = b_idx, fn = fn)
}
}
# phi_<arm> dispersion axes are not block-prefixed; carry them straight
# through so a single `fn_phi` re-weights each on the joint grid.
if (!is.null(fn_phi)) {
for (j in grep("^phi_", cn)) {
out[[length(out) + 1L]] <- list(col = as.integer(j), role = "phi",
block = NA_integer_, fn = fn_phi)
}
}
out
}
# Which axes each regularizing hyperprior reaches, and refusal of a per-block
# entry naming a block that carries no axis of that role. Read off the axis
# layout, so a mis-keyed prior is refused before the first inner solve rather
# than applied to nothing. Returns NULL when no hyperprior is active, else one
# record per role: the scope the caller stated it at, the blocks it reaches and
# the grid columns those are.
.joint_check_multi_hyperpriors <- function(axis_names, axis_offsets, B,
fn_sigma, fn_alpha) {
roles <- list(prior_sigma = "sigma", prior_alpha = "alpha")
fns <- list(prior_sigma = fn_sigma, prior_alpha = fn_alpha)
out <- list()
for (label in names(roles)) {
fn <- fns[[label]]
if (is.null(fn)) next
role <- roles[[label]]
cols <- .joint_block_axis_cols(axis_names, axis_offsets, B, role)
carried <- as.integer(names(cols))
unnamed <- setdiff(.joint_hp_named_blocks(fn), carried)
if (length(unnamed)) {
stop("`", label, "` names block",
if (length(unnamed) > 1L) "s " else " ",
paste(unnamed, collapse = ", "),
", which carr", if (length(unnamed) > 1L) "y" else "ies",
" no `", role, "` axis. ",
if (length(carried))
paste0("The blocks carrying one are ",
paste(carried, collapse = ", "),
" (grid column",
if (length(carried) > 1L) "s " else " ",
paste(axis_names[cols], collapse = ", "), ").")
else paste0("No block of this fit carries a `", role,
"` axis."),
call. = FALSE)
}
applied <- if (is.function(fn)) carried
else intersect(carried, .joint_hp_named_blocks(fn))
out[[role]] <- list(
scope = if (is.function(fn)) "every_block" else "per_block",
blocks = as.integer(applied),
axes = as.character(axis_names[cols[as.character(applied)]]))
}
if (length(out) == 0L) NULL else out
}
# The axis specs a multi-block joint grid's cell measure is built from: the
# specs read off `grid`'s columns with the folded axes marked, the copy scale's
# declared atom mass, and the caller's hyperprior roles
# (`.joint_multi_hp_specs()`), with each free-covariance block measured under
# the design `logchol` declares (`.joint_multi_declared_axes()`). The one
# construction behind the weights and the screen that ranks the cells before any
# is solved.
.joint_multi_measure_specs <- function(grid, folded_axes, joint_grid,
axis_offsets, B, fn_sigma, fn_alpha,
fn_phi, copy_slab, copy_atom_mass,
logchol, axis_refine = NULL) {
specs <- .joint_axis_specs_from_grid(grid, copy_slab = copy_slab,
folded_axes = folded_axes,
logchol = logchol,
axis_refine = axis_refine)
specs <- lapply(specs, function(sp) {
if (!is.null(sp$atom_mass)) sp$atom_mass <- copy_atom_mass
sp
})
.joint_multi_hp_specs(specs, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi)
}
# The hyperprior family of each arm, keyed by arm name, as the default
# hyperprior reads it.
.joint_multi_hp_families <- function(arms, arm_names) {
stats::setNames(
vapply(arms, function(a) as.character(a$family %||% ""), character(1)),
arm_names)
}
# The outer integration of a multi-block joint result whose `log_marginal`
# already carries the folded hyperprior: the hyperprior record over
# `joint_grid`, the cell measure `log_quad` and the axis supports over
# `res$theta_grid` from the specs that record marks, and the integration
# weights, with `dnode` the design weights of a CCD design (NULL on a tensor
# grid). Returns the result and the specs its measure was built from.
#
# `refining` is the per-cell tag the refinement passes leave. The specs are read
# off the grid's BASE cells: an axis's declared node set is what fixes a prior
# read off it, and a node a refinement pass appended was not declared.
.joint_multi_attach_integration <- function(res, joint_grid, axis_offsets, B,
prepared, hp_families,
fn_sigma, fn_alpha, fn_phi,
copy_slab, copy_atom_mass,
dnode = NULL,
hyperprior = "proper", hp_declared,
refining = NULL,
axis_refine = NULL) {
hp <- .joint_multi_hyperprior(
joint_grid, fn_sigma, fn_alpha, fn_phi, blocks = prepared,
families = hp_families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared)
res$log_hyperprior <- hp$lp
res$log_hyperprior_axes <- hp$axes
res$log_hyperprior_declined <- hp$declined
base <- !nzchar(.hyper_slice_home(refining, nrow(res$theta_grid)))
specs <- .joint_multi_measure_specs(
res$theta_grid[base, , drop = FALSE], hp$axes, joint_grid, axis_offsets,
B, fn_sigma, fn_alpha, fn_phi, copy_slab, copy_atom_mass,
logchol = hp_declared$logchol, axis_refine = axis_refine)
res$log_quad <- .hyper_log_quad_weights(res$theta_grid, specs,
refining = refining)
res$axis_support <- .hyper_grid_supports(res$theta_grid, specs,
refining = refining)
res <- .nl_record_unconverged(res)
res$weights <- .joint_integration_weights(res$log_marginal, dnode,
log_quad = res$log_quad)
list(res = res, specs = specs)
}
# Attach each hyperprior to the axis spec of the column it applies to, so the
# axis quadrature knows what the fold below will apply to that axis. That is
# what keeps the copy scale's declared point mass at its declared prior
# probability (`.hyper_atom_fold_scale()`).
.joint_multi_hp_specs <- function(specs, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi = NULL) {
if (is.null(specs)) return(specs)
if (is.null(fn_sigma) && is.null(fn_alpha) && is.null(fn_phi)) return(specs)
entries <- .joint_multi_hp_cols(joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi)
if (length(entries) == 0L) return(specs)
named <- colnames(joint_grid)[vapply(entries, function(e) e[["col"]],
integer(1))]
lapply(specs, function(sp) {
j <- match(sp$name, named)
if (is.na(j)) return(sp)
# The joint driver's hyperprior families are densities on the axis's
# natural scale.
sp$log_prior <- entries[[j]][["fn"]]
sp$log_prior_coord <- "natural"
sp
})
}
# Add the regularizing (sigma, alpha) hyperprior to a per-cell log-marginal
# vector for a multi-block `joint_grid`. Each applied axis is folded through
# the SAME `.joint_hp_vec_for_grids()` the single-block path uses, one column
# at a time, so the atom / continuum reading of a copy scale's zero is that
# axis's own. The per-axis contributions are summed before the one addition to
# `log_marginal`, so a grid carrying a single axis of each role reproduces the
# whole-view call to the bit. Single source of truth for the main dispatch and
# the Pareto-k re-evaluation.
.joint_multi_add_hp <- function(log_marginal, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi = NULL,
blocks = NULL, families = NULL,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
hyperprior = "proper", declared) {
log_marginal + .joint_multi_hyperprior(
joint_grid, fn_sigma, fn_alpha, fn_phi, blocks = blocks,
families = families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = declared)$lp
}
# The hyperprior record over a multi-block `joint_grid` (`.joint_hyperprior()`):
# the caller's densities on the roles they name, and with `blocks` supplied what
# `hyperprior` gives every other axis of each block. `declared` is what the
# fit's declared grid fixes (`.joint_multi_declared_axes()`).
.joint_multi_hyperprior <- function(joint_grid, fn_sigma, fn_alpha, fn_phi = NULL,
blocks = NULL, families = NULL,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
hyperprior = "proper", declared) {
.joint_hyperprior(joint_grid, blocks %||% list(), families,
user = list(sigma = fn_sigma, alpha = fn_alpha,
phi = fn_phi),
copy_atom_mass = copy_atom_mass,
hyperprior = if (is.null(blocks)) "flat" else hyperprior,
declared = declared)
}
# What a multi-block joint grid declares (`.hp_declare()`): the columns it
# integrates -- each block axis declared on more than one node, and each per-arm
# dispersion axis carrying a grid, named as the joint grid names them -- and each
# free-covariance block's design. Read off the declared per-block grids rather
# than any evaluated grid, so the tensor, a CCD design and every batch the driver
# evaluates share one answer.
.joint_multi_declared_axes <- function(block_grids, axis_names,
phi_axes = NULL) {
offsets <- cumsum(c(0L, vapply(block_grids, ncol, integer(1))))
latent <- character(0)
logchol <- list()
for (b in seq_along(block_grids)) {
if (!ncol(block_grids[[b]])) next
g <- block_grids[[b]]
colnames(g) <- as.character(axis_names)[(offsets[b] + 1L):offsets[b + 1L]]
d <- .hp_declare(g)
latent <- c(latent, d$axes)
logchol <- c(logchol, d$logchol)
}
phi <- if (length(phi_axes)) {
active <- vapply(phi_axes, function(v) length(unique(v)) > 1L, logical(1))
paste0("phi_", names(phi_axes)[active])
}
list(axes = c(latent, phi), logchol = logchol)
}
# Multi-block joint outer Pareto-k-hat. Builds the re-evaluation closure
# (round-trips a sampled user-facing `joint_grid` through the SAME kernel via
# the shared cpp-grid / phi / hyperprior helpers) and defers to the
# block-type-aware driver `.joint_pareto_k`. The whole importance batch is
# re-solved in one kernel call across `n_threads_outer`, so the independent
# re-solves run concurrently across cores rather than one-at-a-time using all
# inner threads. No tiling on the re-eval path (the IS batch is not a per-axis
# alpha lattice), so no tile-partition reconstruction is needed.
#
# The diagnostic solves are bounded: each is warm-started
# from the modal cell's converged latent mode (the bulk of the importance
# weight sits near it, so those draws converge in a few Newton steps) and
# capped at `.K_DIAG_MAX_ITER` iterations. A draw at an implausible
# hyperparameter where the inner Newton would otherwise stall to `max_iter`
# carries negligible importance weight, so the cap bounds its cost without
# moving the k-hat (converged draws keep their exact log-marginal; only the
# negligible-weight tail is truncated).
# Fit-scoped factory for the `.cpp_joint_multi` calls on this path (the CCD /
# adaptive mode-find, the node solves, the authoritative solve, the local-CCD
# refinement and the k-diagnostic re-solve). Everything invariant across those
# solves is captured once, so a per-fit setting cannot be honoured at one call
# site and silently dropped at another -- `control$hessian` was: it was parsed
# and match.arg-validated, then never passed on this path at all, so "psd" /
# "fisher" ran as LM/observed without a warning.
#
# The knobs a caller genuinely varies stay parameters. `max_iter` / `tol` /
# `inner_refresh` take fit-scoped defaults but are overridable, because the
# k-diagnostic deliberately re-solves with its own cheaper knobs.
.joint_multi_call_factory <- function(arms, cp, blocks_spec, axis_offsets, B,
n_threads, force_sparse, cell_coupling,
hessian_pd_mode, step_curvature_mode,
max_iter, tol, inner_refresh,
fixed_block_p = 0L,
fixed_block_constraints = NULL,
inner_sparse_override = 0L) {
force(arms); force(cp); force(blocks_spec); force(axis_offsets); force(B)
force(n_threads); force(force_sparse); force(cell_coupling)
force(inner_sparse_override)
force(hessian_pd_mode); force(step_curvature_mode)
force(max_iter); force(tol); force(inner_refresh)
force(fixed_block_p); force(fixed_block_constraints)
function(theta_mat,
x_init = NULL,
store_Q = FALSE,
phi_grid_per_arm = NULL,
n_threads_outer = 1L,
tile_ids = NULL,
tile_pilot_cells = NULL,
prune_tol = 0.0,
x_init_per_cell = NULL,
max_iter_ = NULL,
tol_ = NULL,
inner_refresh_ = NULL,
compute_skew = FALSE,
skew_idx = NULL,
# Per-cell fixed-effect covariance retention.
# Off by default so the mode-find, the Pareto-k re-solve batch and
# the skew probe -- none of which contribute an integrated cell --
# never pay for it; the grid solve and the local-CCD node solves
# switch it on.
fixed_block = FALSE,
# Subspace debias: the kernel-facing request
# list, off by default so no existing call site re-solves anything
# different.
debias = NULL,
# Corrected integrated Laplace: the
# kernel-facing request list, off by default for the same reason.
cila = NULL,
# Per-cell log hyperprior + log cell measure a screened solve
# ranks with, in `theta_mat` row order.
screen_log_offset = NULL,
# Return the screened surface without the full pass: a placement
# pilot's detecting grid only (`.nl_pilot_screen_only()`).
screen_only = FALSE) {
.cpp_joint_multi(
arms_list = arms,
copy_arms = as.integer(cp$copy_arms_zero),
copy_blocks = as.integer(cp$copy_blocks_zero),
blocks_spec = blocks_spec,
theta_grid = .joint_multi_cpp_grid(theta_mat, axis_offsets,
B, cp),
axis_offsets = axis_offsets,
max_iter = as.integer(max_iter_ %||% max_iter),
tol = as.numeric(tol_ %||% 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_ids,
tile_pilot_cells = 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_ %||% inner_refresh),
x_init_per_cell = x_init_per_cell,
compute_skew = isTRUE(compute_skew),
skew_idx = skew_idx,
fixed_block_p = if (isTRUE(fixed_block))
as.integer(fixed_block_p) else 0L,
fixed_block_constraints = fixed_block_constraints,
debias = debias,
cila = cila,
inner_sparse_override = as.integer(inner_sparse_override),
screen_log_offset = screen_log_offset,
screen_only = isTRUE(screen_only))
}
}
# The `kernel_fn(new_cells, warm_start, store_extras)` closure the refinement
# passes drive (`R/hyper_grid_refine.R`), over this driver's own kernel call: the
# multi-block counterpart of `.joint_make_kernel_fn()`. A refinement cell is a row
# of the assembled joint grid, so the per-arm dispersion columns it carries go
# through the same `phi_grid_per_arm` extraction the grid solve uses. The per-cell
# side data leaves only inside `extras`, so a call that asks for none requests
# neither the precision nor the fixed-effect block.
.joint_multi_make_kernel_fn <- function(call_kernel, arm_names, x_init_default,
store_Q = FALSE, screen_tol = 0,
screen_hp = NULL) {
function(new_cells, warm_start = NULL, store_extras = FALSE,
x_init_per_cell = NULL, screen = NULL) {
x_init <- if (!is.null(warm_start) && !is.null(warm_start$mode))
as.numeric(warm_start$mode) else x_init_default
scr <- .nl_refine_screen_offset(new_cells, screen, screen_tol, screen_hp)
res_x <- call_kernel(
new_cells, x_init = x_init,
store_Q = isTRUE(store_Q) && isTRUE(store_extras),
fixed_block = isTRUE(store_extras),
phi_grid_per_arm = .joint_multi_phi_per_arm(new_cells, arm_names),
x_init_per_cell = .joint_warm_start_per_cell(warm_start,
x_init_per_cell,
nrow(new_cells)),
prune_tol = if (is.null(scr)) 0 else screen_tol,
screen_log_offset = scr)
list(log_marginal = res_x$log_marginal,
extras = if (isTRUE(store_extras))
.joint_extras_from_res(res_x, nrow(new_cells)),
modes = res_x$modes)
}
}
# Multi-block counterpart of `.joint_attach_pareto_k_single()`
# (`R/nested_laplace_joint_pareto_k.R`). Same `diagnose_k = FALSE` behaviour:
# `pareto_k` / `pareto_k_is_ess` stay NA, and the placement mode the
# auto-recenter rescues read (`outer_mode_*`, `.joint_attach_placement()`) is
# attached either way.
.joint_attach_pareto_k_multi <- function(res, call_kernel,
axis_offsets, B, arm_names,
fn_sigma, fn_alpha, fn_phi = NULL,
max_iter, tol,
diagnose_k = TRUE,
diagnose_draws = .nl_diag("k_samples"),
n_threads_outer = 1L, proposal = NULL,
pareto_k_by_arm = FALSE,
k_bootstrap = .nl_diag("k_bootstrap"),
k_tail_points = NULL,
k_conf_bands = NULL,
placement_axes = character(0),
hp_blocks = NULL, hp_families = NULL,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
hyperprior = "proper", hp_declared) {
res$pareto_k <- NA_real_
res$pareto_k_is_ess <- NA_real_
res$pareto_k_scope <- "outer (hyperparameter) Gaussian proposal"
res$pareto_k_proposal_source <- NA_character_
res$pareto_k_declined <- NA_character_
# Regime read off stored weights: attached even with the diagnostic off.
res <- .joint_attach_pareto_k_regime(res)
warm_mode <- .joint_modal_mode(res)
modal_theta <- .joint_modal_theta(res)
k_max_iter <- min(as.integer(max_iter), .K_DIAG_MAX_ITER)
n_to <- as.integer(n_threads_outer)
knobs <- .kdiag_knobs()
# One kernel call over the importance batch. Shamanskii reuse makes the
# off-factor steps scatter grad-only (attacks the dominant per-iteration
# Hessian fill); loosened inner tol cuts intrinsic Newton steps; per-cell
# warm start (`x_init_per_cell`, an [S x n_x] matrix or NULL) starts each
# draw from its nearest grid mode in BOTH serial and parallel.
# The diagnostic re-solve runs with its own cheaper knobs (k_max_iter,
# knobs$tol / refresh), so its checkpoint fingerprint would differ from the
# main grid's; run it checkpoint-free so it neither collides with nor
# appends to the fit's checkpoint file.
# Importance draws are laid on each axis's unconstraining transform, the
# identity on a free-covariance block's columns, so they read the density in
# the column coordinates.
k_declared <- .hp_declare_on_columns(hp_declared)
solve_fn <- function(theta_mat, x_init_per_cell = NULL) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat,
x_init = warm_mode,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
n_threads_outer = n_to,
x_init_per_cell = x_init_per_cell,
max_iter_ = k_max_iter,
tol_ = max(knobs$tol, as.numeric(tol)),
inner_refresh_ = knobs$refresh))
.joint_multi_add_hp(r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = hp_blocks, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = k_declared)
}
# The placement mode-find at the fit's own inner settings (see
# `.joint_attach_pareto_k_single()`).
place_warm <- warm_mode
place_fn <- function(theta_mat) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat,
x_init = place_warm,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
n_threads_outer = n_to,
x_init_per_cell = .joint_rows_of(place_warm, nrow(theta_mat))))
lp <- .joint_multi_add_hp(r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = hp_blocks, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = k_declared)
if (is.matrix(r$modes)) attr(lp, "modes") <- r$modes
lp
}
place_set_warm <- function(mode) {
mode <- as.numeric(mode)
if (length(mode) && all(is.finite(mode))) place_warm <<- mode
}
if (!isTRUE(diagnose_k)) {
res <- .k_attach_declined(res, .k_decline("not_requested"))
return(.joint_attach_placement(res, place_fn, place_set_warm,
extra_axes = placement_axes))
}
# Per-cell warm start (nearest grid mode, serial + parallel) when modes are
# stored; else near-neighbour chain re-order; else plain broadcast
refit <- .joint_make_diag_refit(res, solve_fn, modal_theta, knobs)
arm_axes <- if (isTRUE(pareto_k_by_arm)) .joint_pareto_arm_axes(res) else NULL
kd <- .joint_pareto_k(res, refit, diagnose_draws, proposal = proposal,
arm_axes = arm_axes, k_bootstrap = k_bootstrap,
k_tail_points = k_tail_points, k_conf_bands = k_conf_bands)
res$pareto_k <- kd$pareto_k
res$pareto_k_is_ess <- kd$is_ess
res$pareto_k_proposal_source <- kd$proposal_source
res <- .k_attach_declined(res, kd)
res$pareto_k_mode_u <- kd$mode_u
res$pareto_k_cov_u <- kd$cov_u
res$pareto_k_axis_tags <- kd$axis_tags
res$pareto_k_axis_names <- kd$axis_names
res <- .joint_attach_pareto_k_regime(res, kd)
res <- .joint_attach_pareto_k_uncertainty(res, kd)
res <- .joint_attach_by_arm_k(res, kd)
.joint_attach_placement(res, place_fn, place_set_warm,
extra_axes = placement_axes)
}
# Inner-Laplace skewness diagnostic for the multi-block
# joint driver -- the multi-block analogue of
# .nlj_inner_skew_at_theta() (R/nested_laplace_joint.R). Reuses
# `call_kernel` (the SAME
# closure the outer Pareto-k diagnostic and the adaptive-grid refinement
# already call), re-dispatched at the fitted MAP grid cell with
# `compute_skew = TRUE`: one extra deterministic Newton solve, no importance
# sampling -- mirrors the single-block helper exactly except for the
# multi-block `call_kernel` signature (`x_init` instead of `warm_start`, and
# an explicit `phi_grid_per_arm` for the probed cell, since only the
# multi-block path carries per-arm dispersion axes).
#
# Default probe scope is every arm's fixed-effects coefficients, read from
# `res$arm_layout$beta_start` / `$p` (.joint_multi_layout(), already attached
# by the caller) -- the direct multi-block analogue of the single-block
# default ("every arm's beta").
#
# A coupled arm (`cell_coupling != "separable"` on that arm) has no per-obs
# oracle to score (build_joint_curvature3_fns excludes it) -- its indices
# come back NaN, not a silently wrong 0, exactly as for the single-block and
# single-arm paths.
.nlj_multi_inner_skew_at_theta <- function(res, call_kernel, arm_names,
skew_idx = NULL, compute = TRUE) {
res <- .inner_skew_decline(res, "not_requested")
if (!compute) return(res)
modal_theta <- .joint_modal_theta(res)
if (is.null(modal_theta)) return(.inner_skew_decline(res, "backend_unsupported"))
probe_idx <- skew_idx
if (is.null(probe_idx)) {
al <- res$arm_layout
if (is.null(al) || is.null(al$beta_start) || is.null(al$p)) {
return(.inner_skew_decline(res, "no_probe_indices"))
}
probe_idx <- unlist(lapply(seq_len(al$n_arms), function(k) {
if (al$p[k] <= 0L) return(integer(0))
(al$beta_start[k] + 1L):(al$beta_start[k] + al$p[k])
}))
}
probe_idx <- as.integer(probe_idx)
if (length(probe_idx) == 0L) return(.inner_skew_decline(res, "no_probe_indices"))
res <- .inner_skew_decline(res, "solve_failed")
warm_mode <- .joint_modal_mode(res)
theta_mat <- matrix(modal_theta, nrow = 1L,
dimnames = list(NULL, colnames(res$theta_grid)))
out <- tryCatch(
.joint_with_quiet_opts(call_kernel(
theta_mat, x_init = warm_mode,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
compute_skew = TRUE, skew_idx = probe_idx)),
error = function(e) NULL)
.inner_skew_attach_probe(res, out)
}
# Main multi-block joint dispatch. Mirrors the structure of the
# single-block path (build grid, call C++, post-process) but accepts a
# list-of-blocks `prior_list` and a `copy` spec that points at a
# specific block.
.joint_dispatch_multi <- function(responses, prior_list, copy,
phi_grid,
fn_sigma = NULL,
fn_alpha = NULL,
fn_phi = NULL,
max_iter, tol, n_threads,
x_init, verbose, store_Q,
keep_grid_hessians = TRUE,
n_threads_outer = 1L,
tile_warm = TRUE,
prune_tol = 0.0,
force_sparse = FALSE,
inner_sparse_override = 0L,
cell_coupling = "separable",
hessian_pd_mode = 0L,
step_curvature_mode = 0L,
diagnose_k = TRUE,
diagnose_draws = .nl_diag("k_samples"),
pareto_k_by_arm = FALSE,
pareto_k_threads = NULL,
k_bootstrap = .nl_diag("k_bootstrap"),
k_tail_points = NULL,
k_conf_bands = NULL,
diagnose_skew = TRUE,
skew_idx = NULL,
skew_correct = .nl_diag("skew_correct"),
subspace_debias = NULL,
cila = NULL,
inner_refresh = 1L,
integration = "auto",
within_cell = NULL,
local_ccd = NULL,
adaptive_cutoff = 10,
adaptive_stride = 2L,
adaptive_max_frac = 0.75,
adaptive_min_cells = 48,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
copy_slab = "exponential",
placement_axes = character(0),
hyperprior = "proper",
adaptive_grid = FALSE,
adaptive_grid_edge_thresh = 0.02,
adaptive_grid_max_passes = 1L,
var_of_means_consistency = TRUE,
axis_refine = NULL,
grid_auto = list(),
phi_auto = NULL,
timer = NULL) {
tm <- timer %||% .tulpa_timer()
integration <- match.arg(integration, c("auto", "ccd", "grid",
"grid_adaptive"))
within_cell <- .nl_within_cell_mode(within_cell)
n_arms <- length(responses)
arms <- lapply(seq_along(responses), function(k) {
a <- responses[[k]]
.normalise_joint_arm_multi(a, k)
})
arm_names <- names(responses) %||% paste0("arm", seq_along(responses))
cp <- .resolve_copy_multi(copy, responses, prior_list)
hp_families <- .joint_multi_hp_families(arms, arm_names)
# Per-block axis grids (with copy-block parameterisation if applicable).
B <- length(prior_list)
per_block <- lapply(seq_len(B), function(b) {
copy_pos <- if (cp$has_copy) match((b - 1L), cp$copy_blocks_zero)
else NA_integer_
is_copy <- !is.na(copy_pos)
alpha_grid_b <- if (is_copy) cp$alpha_grids[[copy_pos]] else numeric(0)
.joint_block_axis_grid(prior_list[[b]], is_copy, alpha_grid_b, b)
})
block_grids <- lapply(per_block, function(x) x$grid)
prepared <- lapply(per_block, function(x) x$prepared)
# Axis layout (independent of node placement: tensor or CCD).
axis_counts <- vapply(block_grids, ncol, integer(1))
axis_offsets <- as.integer(c(0L, cumsum(axis_counts)))
axis_names <- unlist(lapply(seq_along(block_grids), function(b) {
cn <- colnames(block_grids[[b]])
if (is.null(cn) || length(cn) == 0L) return(character(0))
paste0("b", b, ".", cn)
}))
d_axes <- as.integer(axis_offsets[B + 1L])
# Which axes the regularizing hyperpriors reach. Read off the axis layout,
# before the first inner solve, so a per-block entry naming a block that
# carries no such axis is refused rather than paid for.
hyperprior_axes <- .joint_check_multi_hyperpriors(
as.character(axis_names), axis_offsets, B, fn_sigma, fn_alpha)
if (isTRUE(verbose)) {
for (role in names(hyperprior_axes)) {
rec <- hyperprior_axes[[role]]
if (length(rec$axes) > 1L && identical(rec$scope, "every_block")) {
message("tulpa joint: prior_", role, " applies to ",
paste(rec$axes, collapse = ", "),
" (one spec, every block carrying the axis). ",
"Key it per block to regularize them differently.")
}
}
}
blocks_spec <- lapply(seq_along(prepared), function(b) {
.joint_block_spec_for_cpp(prepared[[b]], n_arms, b, arms = arms)
})
# Every `.cpp_joint_multi` call on this path goes through this closure, so
# the fit-scoped settings are supplied once instead of at each of the seven
# solve sites.
force_sparse <- .resolve_force_sparse(force_sparse, function() {
.joint_multi_layout(arms, prepared)$n_x
})
# Per-cell fixed-effect covariance retention. Both the
# block size and the field sum-to-zero groups are fixed by the latent layout,
# which is known before the first solve, so the request is built here and the
# kernel extracts each cell's block inside that cell's own solve
# instead of the grid keeping every precision.
fixed_layout <- .joint_fixed_layout(responses)
fb_layout <- .joint_multi_layout(arms, prepared)
fixed_block_p <- if (isTRUE(keep_grid_hessians))
as.integer(fixed_layout$n_fixed) else 0L
fixed_block_constraints <- .joint_constraint_cols(fb_layout, fb_layout$n_x)
call_kernel <- .joint_multi_call_factory(
arms = arms, cp = cp, blocks_spec = blocks_spec,
axis_offsets = axis_offsets, B = B, n_threads = n_threads,
force_sparse = force_sparse, cell_coupling = cell_coupling,
hessian_pd_mode = hessian_pd_mode,
step_curvature_mode = step_curvature_mode,
max_iter = max_iter, tol = tol, inner_refresh = inner_refresh,
fixed_block_p = fixed_block_p,
fixed_block_constraints = fixed_block_constraints,
inner_sparse_override = inner_sparse_override)
# Per-latent-axis fine grid values (sorted unique), shared by the CCD and the
# adaptive-lattice integrators to place / locate the outer nodes.
block_of_axis <- rep(seq_len(B), times = axis_counts)
col_within <- unlist(lapply(axis_counts, seq_len))
latent_axis_values <- lapply(seq_len(d_axes), function(j) {
sort(unique(as.numeric(block_grids[[block_of_axis[j]]][, col_within[j]])))
})
# Per-arm dispersion overrides on the outer grid. A phi axis varies
# independently of the latent-block axes, so it forms a Cartesian (tensor)
# product with the latent grid -- whether that latent grid is the dense
# tensor or a CCD design. The cross is applied once,
# below, to whichever latent grid was built.
phi_axes <- .normalise_phi_grid(phi_grid, arm_names)
has_phi <- !is.null(phi_axes) &&
any(vapply(phi_axes, length, integer(1)) > 0L)
hp_declared <- .joint_multi_declared_axes(block_grids, axis_names, phi_axes)
# How far each axis may be refined follows from where its nodes came from,
# the rule the single-block driver applies (`.joint_axis_refine_modes_by()`):
# nodes the caller stated are integrated more finely over exactly that range,
# nodes the engine placed may be followed out past their ends.
# `control$axis_refine` overrides either, per axis, and is resolved before
# the first kernel call so a mis-named axis is refused rather than reported
# after the grid has been solved.
refine_axes <- c(as.character(axis_names),
if (has_phi) paste0("phi_", names(phi_axes)[
vapply(phi_axes, length, integer(1)) > 0L]))
axis_refine_modes <- .joint_axis_refine_modes_by(
refine_axes,
function(a) .joint_multi_axis_is_stated(
a, prior_list, grid_auto, copy, as.integer(cp$copy_blocks_zero) + 1L,
phi_grid, arm_names, phi_auto),
user = .joint_check_axis_refine(axis_refine, refine_axes))
# Node placement for the LATENT axes. CCD integrates on a
# central composite design around the joint hyperparameter mode; it declines
# (-> tensor grid) for too few axes, an unguessable axis (CAR_proper rho_car
# / non-BYM2 rho), or a degenerate / ridged outer mode-find / Hessian
#. An active phi axis no longer disables CCD
#: the CCD rides the latent axes and the phi tensor
# crosses on top, with the CCD design weights replicated across phi cells.
# Axis threshold: "auto" (the default) engages CCD only at d >= 4, where the
# tensor blow-up bites hardest, and keeps the cheaper, ridge-robust tensor
# grid at d <= 3; explicit "ccd" lowers the threshold to d >= 3.
dnode <- NULL
tile_partition <- NULL
phi_grid_per_arm_list <- NULL
integration_used <- "grid"
joint_grid <- NULL
use_adaptive <- FALSE
adaptive_declined <- FALSE
adaptive_info <- NULL
# Mode-Hessian outer proposal. Set only when the CCD
# integrator engages: its Gaussian (centre `u_hat`, scale `L_scale` from the
# analytic curvature at the outer mode) drives the outer Pareto-k so the
# diagnostic stays defined when the hyperparameter posterior is sharp and the
# grid concentrates on ~1 cell (where the grid-weighted covariance collapses).
ccd_proposal <- NULL
ccd_requested <- .joint_ccd_engage(integration, d_axes)
# Why a CCD the caller asked for did not run. NA when
# nothing was declined, so `fit$integration == "grid"` is never the only
# trace of a request the engine turned down. An explicit
# `integration = "ccd"` below the axis threshold declines here; every other
# reason comes from .joint_ccd_grid()'s own vocabulary.
integration_declined <- NA_character_
if (identical(integration, "ccd") && !ccd_requested) {
integration_declined <- "axis_count"
}
# What placing the design cost, in inner solves. Carried on the placement's
# own return so a decline reports what it had already paid, and NULL where
# no placement ran -- a fit that placed nothing reports no cost rather than
# a zero, which would read as a placement that was free.
ccd_cost <- NULL
use_ccd <- ccd_requested
# Set only on the Cartesian-product branch below (CCD and the adaptive
# grid do not build a dense tensor, so the "reduce the grid" remedy the
# post-solve timing check reports does not apply to them).
n_cells <- NA_integer_
grid_warn_remedy <- NULL
grid_layout <- ""
if (use_ccd) {
axis_values <- latent_axis_values
# Warm latent mode for the mode-find / node solves: one solve at the
# box-centre coordinate, broadcast as x_init so each subsequent inner
# Newton converges in a few steps. Checkpoint/progress are stripped for
# every CCD probe so they neither pollute the fit's checkpoint file nor
# tick its progress bar.
ccd_warm <- x_init
centre_row <- matrix(vapply(axis_values, stats::median, numeric(1)),
nrow = 1L, dimnames = list(NULL, axis_names))
init_fit <- tryCatch(.joint_with_quiet_opts(call_kernel(
centre_row, x_init = x_init)),
error = function(e) NULL)
if (!is.null(init_fit) && is.matrix(init_fit$modes) &&
nrow(init_fit$modes) >= 1L &&
is.finite(init_fit$log_marginal[1L])) {
m1 <- as.numeric(init_fit$modes[1L, ])
if (all(is.finite(m1))) ccd_warm <- m1
}
# A CCD design is laid, and its mode found, in the grid's column
# coordinates, so the hyperprior it reads is the density in those.
ccd_declared <- .hp_declare_on_columns(hp_declared)
eval_logpost <- function(theta_mat) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat, x_init = ccd_warm,
n_threads_outer = n_threads_outer))
lp <- .joint_multi_add_hp(r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = prepared, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior,
declared = ccd_declared)
# Carry the inner latent modes so the CCD mode-find can advance the
# warm start per accepted point.
if (is.matrix(r$modes)) attr(lp, "modes") <- r$modes
lp
}
# Advance the inner warm start to the latent mode at each accepted
# mode-find point, so subsequent probes solve in a few Newton steps
# rather than cold from the box centre.
set_warm <- function(mode) {
mode <- as.numeric(mode)
if (length(mode) && all(is.finite(mode))) ccd_warm <<- mode
}
# The same atom mass the axis specs below declare, so the mixture the
# design splits into is weighted by the prior the grid is integrated
# against.
ccd <- .joint_ccd_grid(axis_names, axis_offsets, prepared, axis_values,
eval_logpost, verbose = verbose,
set_warm = set_warm,
atom_mass = copy_atom_mass)
ccd_cost <- ccd[c("modefind_evals", "modefind_budget",
"modefind_rounds", "modefind_seconds")]
if (is.null(ccd$grid)) {
use_ccd <- FALSE
integration_declined <- as.character(ccd$declined %||% "modefind_failed")
} else {
joint_grid <- ccd$grid
dnode <- ccd$dnode
integration_used <- "ccd"
hp_declared <- ccd_declared
# The CCD axes are exactly the leading latent-block columns of
# joint_grid (axis_names, in order); phi crosses as a separate
# tensor on top. Carry the mode-Hessian Gaussian over those axes for
# the outer Pareto-k.
#
# A design that split at a copy atom has no such Gaussian to carry.
# Its components are built in log alpha, while the diagnostic
# unconstrains the SAME grid on the identity coordinate -- the only
# one that represents the alpha = 0 cells the split places. One
# Gaussian does not describe a mixture either way, so the diagnostic
# keeps its grid-weighted proposal, which the split's spread of
# components supports.
if (isTRUE(ccd$atom_split)) {
ccd_proposal <- NULL
} else {
ccd_proposal <- list(u_hat = ccd$u_hat,
L_scale = ccd$L_scale,
tags = ccd$tags,
cols = seq_len(d_axes))
}
}
}
# Adaptive lattice (nested_laplace_joint_adaptive.R): the low-dimensional
# companion to the CCD. It evaluates a coarse subsample of the full outer
# lattice (latent block axes AND phi axes together) to locate the posterior
# mass, then floods outward on the fine lattice keeping only the cells within
# a log-density cutoff of the peak -- so the far tail, which every dense cell
# still pays a full inner solve for, is never placed. The selected cells are
# a strict subset of the dense tensor at uniform lattice spacing, so the main
# kernel call below evaluates them with the full contract (store_Q, phi, tile
# warm start) and the integration weight stays the plain softmax (dnode NULL);
# the flood solves here are cheap grid selection (store_Q off, quiet). Declines
# to the dense tensor on a degenerate lattice or a diffuse posterior whose kept
# region rivals the dense grid.
if (!use_ccd && .joint_adaptive_engage(integration, d_axes)) {
if (has_phi) {
active_ad <- phi_axes[vapply(phi_axes, length, integer(1)) > 0L]
phi_axis_vals <- lapply(active_ad, function(v)
sort(unique(as.numeric(v))))
phi_cols_ad <- paste0("phi_", names(active_ad))
} else {
phi_axis_vals <- list()
phi_cols_ad <- character(0)
}
ad_axis_values <- c(latent_axis_values, phi_axis_vals)
ad_col_names <- c(axis_names, phi_cols_ad)
# One warm centre solve, broadcast as x_init so each flood-selection solve
# converges in a few Newton steps rather than cold from the box centre.
ad_warm <- x_init
centre_theta <- matrix(
vapply(ad_axis_values, stats::median, numeric(1)), nrow = 1L,
dimnames = list(NULL, ad_col_names))
init_ad <- tryCatch(.joint_with_quiet_opts(call_kernel(
centre_theta, x_init = x_init,
phi_grid_per_arm = .joint_multi_phi_per_arm(centre_theta, arm_names))),
error = function(e) NULL)
if (!is.null(init_ad) && is.matrix(init_ad$modes) &&
nrow(init_ad$modes) >= 1L && is.finite(init_ad$log_marginal[1L])) {
m1 <- as.numeric(init_ad$modes[1L, ])
if (all(is.finite(m1))) ad_warm <- m1
}
# Grid-selection evaluator: physical theta (latent + phi columns) -> the
# inner joint log-marginal with the hyperprior baked in. store_Q off and
# quiet, since these solves only rank cells; the authoritative solve (with
# store_Q) is the main kernel call over the selected grid.
eval_theta_ad <- function(theta_mat) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat, x_init = ad_warm,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
n_threads_outer = n_threads_outer))
list(log_marginal = .joint_multi_add_hp(
r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = prepared, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared),
modes = NULL)
}
ad <- .joint_adaptive_grid(ad_axis_values, ad_col_names, eval_theta_ad,
cutoff = adaptive_cutoff,
stride = adaptive_stride,
max_frac = adaptive_max_frac,
min_dense = adaptive_min_cells,
verbose = verbose)
adaptive_declined <- !is.null(ad$declined)
if (adaptive_declined) {
integration_declined <- ad$declined
} else {
joint_grid <- ad$grid
integration_used <- "grid_adaptive"
use_adaptive <- TRUE
adaptive_info <- ad$info
# phi columns are already part of the adaptive grid; extract them for
# the C++ call and skip the phi tensor cross below.
phi_grid_per_arm_list <- .joint_multi_phi_per_arm(joint_grid, arm_names)
}
}
if (!use_ccd && !use_adaptive) {
# Cartesian product of per-block axis grids.
# The product over zero blocks is one cell with no latent axis, so a
# model whose only hyperparameters are dispersions is solved once per
# dispersion cell crossed on top.
row_counts <- vapply(block_grids, nrow, integer(1))
idx <- if (B > 0L) do.call(expand.grid, lapply(row_counts, seq_len))
else data.frame(row.names = 1L)
# Every latent cell is solved once per dispersion cell crossed on top
# below, so the count the timing warning reads is the crossed one: that
# is the number of inner solves the fit pays.
phi_active <- if (has_phi) phi_axes[lengths(phi_axes) > 0L]
n_cells <- nrow(idx) * prod(lengths(phi_active))
# The CCD and the adaptive lattice are remedies for the cell count only
# while they are still available. Once this fit has already declined
# one, advising the caller to set the option they set is noise, so the
# advice names the decline instead. The adaptive lattice is named first:
# it integrates the same lattice, where the CCD is a different design.
alt <- c(
if (!adaptive_declined)
"\"grid_adaptive\" (the same lattice, solved where the posterior mass sits)",
if (adaptive_declined || is.na(integration_declined)) "\"ccd\"")
declined <- if (!is.na(integration_declined))
sprintf(if (adaptive_declined)
"the adaptive lattice declined to the dense tensor (%s)"
else "CCD integration declined (%s)",
integration_declined)
grid_warn_remedy <- paste0(
if (length(phi_active)) "Reduce per-block or phi grid sizes"
else "Reduce per-block grid sizes",
if (length(alt))
paste0(" or set control$integration = ",
paste(alt, collapse = " or ")),
if (length(declined))
paste0("; ", paste(declined, collapse = "; ")),
".")
grid_layout <- .nl_grid_crossing(block_grids, phi_active)
.nl_check_grid_cap(n_cells, .nl_max_grid_cells(), grid_warn_remedy,
grid_layout)
.nl_dense_grid_warn(n_cells, grid_warn_remedy, grid_layout)
joint_grid <- if (B > 0L) do.call(cbind, lapply(seq_along(block_grids),
function(b) block_grids[[b]][idx[[b]], , drop = FALSE]))
else matrix(numeric(0), nrow = 1L, ncol = 0L)
if (ncol(joint_grid) > 0L) colnames(joint_grid) <- axis_names
}
# Latent node / cell count, before the phi tensor crosses on top -- the CCD
# node count or the tensor latent-cell count, for the verbose announcement.
n_latent_cells <- nrow(joint_grid)
# phi_grid (per-arm dispersion overrides on the outer grid): each phi axis
# varies independently of the latent-block grid, so it forms a Cartesian
# product with the latent grid -- the dense tensor or the CCD design alike
#. On the CCD path the per-node design weights `dnode`
# are replicated across the phi cells so each (latent node, phi cell) keeps
# the latent node's quadrature weight; phi rides as a uniform tensor axis,
# exactly as on the dense path. A phi axis tied to a latent-block axis is
# not expressible here; fold it into the relevant block's own axis grid.
if (has_phi && !use_adaptive) {
crossed <- .joint_multi_cross_phi(joint_grid, phi_axes, axis_names)
joint_grid <- crossed$grid
# phi columns don't belong to any latent block; the C++ entry consumes
# them via phi_grid_per_arm, not the block axes.
phi_grid_per_arm_list <- .joint_multi_phi_per_arm(joint_grid, arm_names)
# Replicate the CCD design weights across the crossed phi cells (NULL on
# the dense path, where weights are the uniform-cell softmax).
if (!is.null(dnode)) dnode <- dnode[crossed$latent_row]
}
# Announce the engaged outer integrator at selection time:
# the single authoritative line, tied to integration_used, before the heavy
# inner solves -- so a CCD engaged silently by axis count (or declined back
# to the tensor on a ridge) is visible without reading fit$...$integration.
if (isTRUE(verbose) && !use_adaptive) {
n_phi_cells <- as.integer(round(nrow(joint_grid) / n_latent_cells))
.joint_announce_integration(
integration_used, d_axes, n_latent_cells, n_phi_cells,
nrow(joint_grid),
declined = if (adaptive_declined) "adaptive lattice"
else if (ccd_requested && !identical(integration_used, "ccd"))
"CCD")
}
# Tile partition for the three-tier warm-start.
# Tensor-grid-only: the scattered CCD design is not a
# per-axis lattice the alpha-tile reuse relies on. Tile axis = every
# joint_grid column EXCEPT the copy block's alpha column. Built from the
# *user-facing* (sigma, alpha) grid (before sigma_pos materialisation) so
# the partition reflects what is constant across alpha at fixed (sigma, rho,
# ... other-block axes).
if (!use_ccd && isTRUE(tile_warm) && cp$has_copy &&
length(cp$copy_blocks_zero) == 1L &&
as.integer(n_threads_outer) > 1L) {
b_copy_R <- cp$copy_blocks_zero[[1L]] + 1L
cols_bc <- (axis_offsets[b_copy_R] + 1L):axis_offsets[b_copy_R + 1L]
bare_bc <- sub("^b[0-9]+\\.", "", colnames(joint_grid)[cols_bc])
i_alpha <- match("alpha", bare_bc)
if (!is.na(i_alpha)) {
alpha_col_idx <- cols_bc[i_alpha]
non_alpha_idx <- setdiff(seq_len(ncol(joint_grid)), alpha_col_idx)
tile_partition <- .joint_compute_tile_partition(
joint_grid[, non_alpha_idx, drop = FALSE],
as.numeric(joint_grid[, alpha_col_idx]),
nrow(joint_grid)
)
}
}
# The regularizing hyperprior, folded into log_marginal at the kernel-call
# boundary below, and the base tensor's cell measure: known from the grid
# alone, so a screened solve ranks its cells with both.
hp_base <- .joint_multi_hyperprior(
joint_grid, fn_sigma, fn_alpha, fn_phi, blocks = prepared,
families = hp_families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared)
screen_offset <- if (as.numeric(prune_tol) > 0)
.nl_screen_log_offset(joint_grid, list(hp_base),
specs = .joint_multi_measure_specs(
joint_grid, hp_base$axes, joint_grid,
axis_offsets, B, fn_sigma, fn_alpha, fn_phi,
copy_slab, copy_atom_mass,
logchol = hp_declared$logchol))
# A placement pilot's detecting grid is read off the screen alone
# (`.nl_pilot_control()`); the full-grid fallback below never is.
pilot_screen <- isTRUE(getOption("tulpa.nl_pilot_screen", FALSE))
call_kernel_with_tol <- function(tol_prune) {
call_kernel(
joint_grid,
x_init = x_init,
store_Q = store_Q,
fixed_block = TRUE,
phi_grid_per_arm = phi_grid_per_arm_list,
n_threads_outer = n_threads_outer,
tile_ids = tile_partition$tile_ids,
tile_pilot_cells = tile_partition$tile_pilot_cells,
prune_tol = tol_prune,
screen_log_offset = if (tol_prune > 0) screen_offset,
screen_only = pilot_screen && tol_prune > 0)
}
tm$mark("setup")
grid_solve_start <- proc.time()[["elapsed"]]
res <- .joint_main_grid_solve(function() call_kernel_with_tol(prune_tol))
# Safety gate: fall back to the full grid when the
# cheap-pass ranking is unreliable rather than silently returning a
# pruned answer.
if (as.numeric(prune_tol) > 0) {
res <- .joint_prune_safety_gate(
res, resolve_full = function() call_kernel_with_tol(0.0))
}
if (!is.null(grid_warn_remedy)) {
.nl_multi_grid_warn(proc.time()[["elapsed"]] - grid_solve_start,
n_cells, grid_warn_remedy, grid_layout)
}
tm$mark("grid")
# Bake the regularizing hyperprior on (sigma, alpha) into log_marginal
#. Multi-block has no in-package refinement passes,
# so one apply at the kernel-call boundary suffices.
res <- .nl_fold_hyperprior(res, list(hp_base))
# The base tensor's absolute cell measure, read before local CCD refinement
# replaces cells by design-weighted clouds whose shares are relative to it.
base_grid <- joint_grid
# Local CCD refinement: replace a few high-weight, mutually
# non-adjacent tensor cells with small curvature-aware node clouds so a coarse
# base grid resolves the sharply-peaked directions without the k^d tensor
# blow-up. Engages only on the tensor path (the finite-difference curvature
# stencil needs axis neighbours), at >= 4 transformable latent axes, with no
# active phi grid. The node solves reuse the kernel through a quiet evaluator
# (checkpoint / progress stripped) warm-started per cell; refined cells carry
# partition-of-unity design weights so the total integration weight is
# conserved (no double-count).
local_ccd_info <- NULL
weight_kind <- NULL
# Local-CCD refinement changes the outer grid but not the stored per-cell Q
# (nor the tile partition / phi tensor), so combining it with store_Q would
# leave a Q that no longer aligns with the refined grid and crash
# tulpa_posterior_draws. Keeping Q valid for draws takes precedence: skip
# the refinement when store_Q is on.
if (!is.null(local_ccd) && isTRUE(store_Q) && !use_ccd && !use_adaptive &&
!has_phi && .joint_local_ccd_engage(d_axes) && isTRUE(verbose)) {
message("tulpa joint: local-CCD refinement skipped (store_Q = TRUE); ",
"keeping the per-cell Q aligned with the grid for posterior draws.")
}
if (!is.null(local_ccd) && !isTRUE(store_Q) && !use_ccd && !use_adaptive &&
!has_phi && .joint_local_ccd_engage(d_axes)) {
lc <- if (is.list(local_ccd)) local_ccd else list()
lc_max_cells <- as.integer(lc$max_cells %||% 8L)
lc_f0 <- as.numeric(lc$f0 %||% 1.1)
lc_skew_max <- as.numeric(lc$skew_max %||% .nl_diag("gamma3_ok"))
lc_rank <- match.arg(lc$rank %||% "weight", .LCCD_RANKS)
tags_lc <- .joint_ccd_axis_tags(axis_names, axis_offsets, prepared)
if (!is.null(tags_lc)) {
eval_nodes <- function(theta_mat, warm) {
xi <- if (!is.null(warm) && length(warm) && all(is.finite(warm)))
as.numeric(warm) else x_init
r <- .joint_with_quiet_opts(call_kernel(
theta_mat, x_init = xi,
fixed_block = (fixed_block_p > 0L),
n_threads_outer = n_threads_outer))
list(log_marginal = .joint_multi_add_hp(
r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = prepared, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared),
modes = if (is.matrix(r$modes)) r$modes else NULL,
cov_blocks = r$cov_block_per_grid)
}
ref <- .joint_local_ccd_refine(
joint_grid = joint_grid, log_marginal = res$log_marginal,
modes = res$modes, cov_blocks = res$cov_block_per_grid,
dnode = dnode, latent_axes = axis_names,
tags = tags_lc, eval_nodes = eval_nodes,
max_cells = lc_max_cells, f0 = lc_f0, verbose = verbose,
skew_max = lc_skew_max, rank = lc_rank)
if (!is.null(ref)) {
joint_grid <- ref$joint_grid
res$log_marginal <- ref$log_marginal
res$modes <- ref$modes
dnode <- ref$dnode
local_ccd_info <- ref$info
weight_kind <- ref$weight_kind
# The node solves carry their own fixed-effect block, so the
# refined grid keeps the retention aligned with the weights it
# is integrated against. The per-cell
# precision is not carried, which is why refinement engages only
# when store_Q is off -- the gate above.
res$cov_block_per_grid <- ref$cov_blocks
# The per-cell linear predictor is not carried either, and
# unlike the precision it is not gated: it still indexes the
# grid the fit had BEFORE this pass. Dropped rather than read
# against the new cells, so the grid-mixture predictive read
# falls back to the fit's coefficients instead of drawing from
# cells that are no longer the fit's (gcol33/tulpa#850).
res$fitted_eta <- NULL
res$fitted_eta_var <- NULL
}
}
tm$mark("grid")
}
# Outer-grid refinement on the tensor grid, the passes the single-block
# driver runs (`.joint_refine_outer_grid()`): an axis whose marginal has
# collapsed onto too few nodes to carry a spread is bisected in the modal
# cell's row until it resolves, and the boundary / interior pass follows the
# posterior past an engine-placed axis's ends when asked for. A CCD design,
# the adaptive lattice and a locally refined grid are not lattices a node can
# be inserted into, so they keep the nodes they were laid on.
joint_grid_init <- joint_grid
refining_axis <- NULL
refine_on <- !use_ccd && !use_adaptive && is.null(local_ccd_info) &&
(isTRUE(adaptive_grid) || isTRUE(var_of_means_consistency))
refine_hp_record <- function(cells)
.joint_multi_hyperprior(
cells, fn_sigma, fn_alpha, fn_phi, blocks = prepared,
families = hp_families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared)
refine_hp <- function(cells) refine_hp_record(cells)$lp
refine_kernel <- .joint_multi_make_kernel_fn(call_kernel, arm_names, x_init,
store_Q,
screen_tol = as.numeric(prune_tol),
screen_hp = refine_hp_record)
grid_lm <- res$log_marginal
grid_extras <- .joint_init_extras_from_res(res)
n_added <- 0L
if (refine_on) {
refine_specs <- .joint_multi_measure_specs(
joint_grid, hp_base$axes, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi, copy_slab, copy_atom_mass,
logchol = hp_declared$logchol, axis_refine = axis_refine_modes)
ref <- .joint_refine_outer_grid(
joint_grid, grid_lm, grid_extras, refine_specs, refine_kernel,
refine_hp,
adaptive_grid = adaptive_grid,
edge_thresh = adaptive_grid_edge_thresh,
max_passes = adaptive_grid_max_passes,
consistency = var_of_means_consistency,
axis_modes = .nl_outer_mode_axes(
getOption("tulpa.nl_outer_mode")))
refining_axis <- ref$refining_axis
joint_grid <- ref$theta_grid
grid_lm <- ref$log_marginal
grid_extras <- ref$extras
n_added <- ref$n_added
adaptive_info <- ref$adaptive_info
if (isTRUE(var_of_means_consistency)) {
res$var_of_means_consistency_info <- ref$consistency_info
}
}
# A cell whose inner Newton stopped short of a mode, on any node set this
# driver lays, is solved again from its nearest converged neighbour's mode
# before the grid is read.
rs <- .joint_resolve_unconverged(joint_grid, grid_lm, grid_extras,
refine_kernel, refine_hp)
if (n_added > 0L || rs$info$n_resolved > 0L) {
res <- .joint_glue_extras_to_res(res, joint_grid, rs$log_marginal,
rs$extras, refining_axis)
}
res$refining_axis <- refining_axis
res$nonconverged_resolve <- rs$info
tm$mark("grid")
res$theta_grid <- joint_grid
res$theta_names <- colnames(joint_grid)
res$axis_offsets <- axis_offsets
# Which grid axes each regularizing hyperprior was folded onto, and whether
# the caller stated it once (every block carrying that axis) or per block.
# NULL where the fit carries none, so a flat fit reports nothing rather
# than an empty record.
res$hyperprior_axes <- hyperprior_axes
res$integration <- integration_used
# What the caller asked for, and why it did not run. `integration` alone
# names the integrator that RAN, so a fit reporting "grid" cannot say
# whether the tensor grid was chosen or fallen back to.
# `.nl_node_support()` keys the interval construction off `integration`, so
# the difference is one a consumer can act on.
res$integration_requested <- integration
res$integration_declined <- integration_declined
res$ccd_modefind_evals <- ccd_cost$modefind_evals
res$ccd_modefind_budget <- ccd_cost$modefind_budget
res$ccd_modefind_rounds <- ccd_cost$modefind_rounds
res$ccd_modefind_seconds <- ccd_cost$modefind_seconds
# Integration weights fold in the CCD design weights (`dnode`) and, on the
# tensor grid, the per-cell prior mass of the node each cell sits at. A CCD
# design carries its own volume in `dnode`, so `log_quad` applies to the
# tensor path alone.
integ <- .joint_multi_attach_integration(
res, joint_grid, axis_offsets, B, prepared, hp_families,
fn_sigma, fn_alpha, fn_phi, copy_slab, copy_atom_mass, dnode = dnode,
hyperprior = hyperprior, hp_declared = hp_declared,
refining = refining_axis,
axis_refine = if (refine_on) axis_refine_modes)
res <- .nl_prune_merged_record(integ$res)
multi_specs <- integ$specs
if (refine_on) {
res$axis_span <- .joint_axis_span(joint_grid_init, res$theta_grid,
multi_specs, refining = refining_axis)
}
# The outer design weight each cell carries, kept beside the integration
# weight it was folded into: absent on a tensor base (uniform cell weight),
# the CCD design weights on a global CCD, the partition-of-unity shares of
# each replacement cloud on a locally refined grid. Recovering it from
# `weights` afterwards is a division by exp(log_marginal), which loses its
# scale and is undefined on a cell whose inner solve returned no finite
# marginal, so the grid state a weight-construction experiment reads is
# recorded rather than reconstructed.
res$dnode <- dnode
is_ccd <- identical(integration_used, "ccd")
# The evidence is read under the same measure the weights take: the cell
# masses on a tensor grid, the partition-of-unity design shares on a
# locally refined one. A global CCD's design weights reproduce moments and
# carry no volume, so no evidence is read off them.
if (is_ccd) {
res$log_evidence <- NA_real_
res$log_evidence_declined <- "moment_rule_design"
} else if (is.null(dnode)) {
res <- .nl_attach_evidence(res, res$theta_grid, multi_specs)
} else {
# A refined cloud's design shares split the base cell it replaced, so
# each node's absolute measure is its share of that cell's. The base
# tensor is evenly spaced wherever refinement engages, so one cell's
# measure serves every node; an uneven base has no single value to share.
base_specs <- .joint_axis_specs_from_grid(
base_grid, copy_slab = copy_slab,
folded_axes = res$log_hyperprior_axes)
base_lq <- .hyper_log_quad_weights(base_grid, base_specs, absolute = TRUE)
even <- all(is.finite(base_lq)) && diff(range(base_lq)) < 1e-8
res <- .nl_attach_evidence(
res, res$theta_grid, multi_specs,
log_measure = if (even) log(pmax(dnode, 0)) + base_lq[1L]
else rep(NA_real_, length(dnode)))
}
# What kind of weight each cell carries. `integration` names the integrator
# that ran, which describes a homogeneous support: a tensor cell holds the
# mass of its own cell, a CCD node holds a design weight. A locally refined
# grid is the one case with both, so it says so per cell rather than leaving
# a consumer to read one kind off `integration`.
res$weight_kind <- weight_kind %||%
rep(if (is_ccd) "design" else "mass", nrow(joint_grid))
if (!is.null(local_ccd_info)) {
# The share of the posterior that sits on design-weighted nodes: the
# part of the support whose cumulative weight is not a CDF, and so the
# bound on how much mass the per-axis quantile can misplace. It is also
# confined to the refined cells' own Voronoi boxes, so the displacement
# is at most one base-grid cell wide.
local_ccd_info$design_mass <-
sum(res$weights[res$weight_kind == "design"], na.rm = TRUE)
}
# Local CCD refinement also leaves a design-weighted (scattered) grid, so it
# takes the same moment / SD path as the global CCD: weighted moments over the
# design, no per-axis lattice SD refit. The adaptive lattice is a mass-
# concentrated SUBSET of the tensor (gaps in the far tail), so its per-axis
# levels are not the regular full lattice the 3-point profile SD refit needs;
# and it is dense near the peak, so the weighted var-of-means over the kept
# cells is already the calibrated SD. It therefore takes the same weighted-
# moment path, not the lattice refit.
# A refinement whose every candidate cell was put back as a mass atom
# leaves a plain tensor grid, so it takes the density path
# its cells carry rather than the design-weighted one it asked for.
is_design_weighted <- is_ccd || use_adaptive ||
(!is.null(local_ccd_info) && local_ccd_info$n_design_nodes > 0L)
# The median and interval per axis are read off whatever node set the
# integrator left. The global CCD is a MOMENT RULE -- its nodes reproduce the
# integrand's moments and carry no mass -- so a cumulative sum over them is
# not a CDF and the interval comes from the moments instead. The adaptive
# lattice keeps a cell per unit of mass, so its cumulative weight is a CDF.
# A locally CCD-refined grid is neither: it is part cell masses and part
# design, and says so rather than passing for a density grid. What it
# REPORTS is still the weighted quantile -- that measured best against a
# converged reference in both design_mass regimes, see
# `.nl_summary_quantile` -- so naming the support changes what the fit can
# tell a caller, not the numbers.
prov <- .nl_interval_provenance(integration_used, res$weight_kind,
res$weights)
res$theta_interval_read <- prov$read
res$theta_interval_design_mass <- prov$design_mass
res <- .joint_posterior_moments_multi(
res, prepared, axis_offsets, joint_grid, cp,
int_weights = if (is_design_weighted) res$weights else NULL,
support = prov$read, within = within_cell)
res$arm_layout <- .joint_multi_layout(arms, prepared)
res$blocks <- prepared
res$prior <- prior_list
res$responses <- responses
res <- .joint_attach_observations(res, responses)
res$copy <- copy
res$cell_coupling <- cell_coupling
res$local_ccd_info <- local_ccd_info
res$adaptive_grid_info <- adaptive_info
# Every axis's declared nodes, whatever subset of the lattice the adaptive
# flood, a pilot thinning or the prune screen kept: the placement mode-find
# searches an axis over these (`.nl_placement_mode()`).
phi_declared <- if (has_phi)
phi_axes[vapply(phi_axes, length, integer(1)) > 0L] else list()
res$outer_axis_nodes <- stats::setNames(
c(latent_axis_values,
lapply(phi_declared, function(v) sort(unique(as.numeric(v))))),
c(axis_names, if (length(phi_declared))
paste0("phi_", names(phi_declared))))
tm$mark("postproc")
# Outer Pareto-k-hat: re-evaluate the inner joint marginal at sampled
# hyperparameters via the shared cpp-grid / phi / hyperprior helpers and
# PSIS-smooth. Declines (NA -> quad-ESS) when a block carries an axis with
# unguessable support (CAR_proper's rho_car, a non-BYM2 rho, ...).
# The diagnostic's importance batch runs across `pareto_k_threads` (resolved
# by the top-level caller from control$k_threads; NULL only on a direct call,
# where it follows this fit's own thread grant)..
k_to <- pareto_k_threads %||%
.tulpa_pareto_k_threads(n_threads_outer, n_threads, diagnose_draws, NULL)
res <- .joint_attach_pareto_k_multi(res, call_kernel,
axis_offsets, B, arm_names,
fn_sigma, fn_alpha, fn_phi,
max_iter, tol,
diagnose_k = diagnose_k,
diagnose_draws = diagnose_draws,
n_threads_outer = k_to,
proposal = ccd_proposal,
pareto_k_by_arm = pareto_k_by_arm,
k_bootstrap = k_bootstrap,
k_tail_points = k_tail_points,
k_conf_bands = k_conf_bands,
placement_axes = placement_axes,
hp_blocks = prepared,
hp_families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior,
hp_declared = hp_declared)
res <- .nlj_multi_inner_skew_at_theta(res, call_kernel, arm_names,
skew_idx, compute = diagnose_skew)
fixed <- .joint_fixed_layout(responses)
res <- .nl_skew_correction_attach(res, fixed$n_fixed, skew_correct)
tm$mark("diagnostics")
# Per-cell fixed-effect mode + precision for the grid marginalization.
# Local-CCD refinement invalidates the stored cells, so a fit it engaged on
# records that reason instead.
res <- .joint_finalize_grid_fixed(res, fixed$n_fixed, keep_grid_hessians)
# Subspace debias, over the settled grid, through the
# same `call_kernel` the skew probe and the Pareto-k batch re-dispatch with.
res <- .nl_subspace_debias_attach(
res, subspace_debias,
redispatch = function(req) .joint_with_quiet_opts(call_kernel(
res$theta_grid,
phi_grid_per_arm = .joint_multi_phi_per_arm(res$theta_grid,
arm_names),
debias = req)),
p_fixed = fixed$n_fixed, beta_names = fixed$names)
# Corrected integrated Laplace, over the same settled
# grid and through the same `call_kernel`.
res <- .nl_cila_attach(
res, cila,
redispatch = function(req) .joint_with_quiet_opts(call_kernel(
res$theta_grid,
phi_grid_per_arm = .joint_multi_phi_per_arm(res$theta_grid,
arm_names),
cila = req)),
p_fixed = fixed$n_fixed, beta_names = fixed$names,
remoments = function(r) {
r <- .joint_posterior_moments_multi(
r, prepared, axis_offsets, joint_grid, cp,
int_weights = if (is_design_weighted) r$weights else NULL,
support = prov$read, within = within_cell)
r
})
res$timing <- tm$timing()
res <- .joint_attach_diagnose_cost(res, diagnose_k, diagnose_draws)
.finalize_fit(res, backend = "nested_laplace_joint",
n_fixed = fixed$n_fixed, fixed_names = fixed$names,
extra_class = c("tulpa_nested_laplace_joint_multi",
"tulpa_nested_laplace_joint",
"tulpa_nested_laplace", "list"))
}
# Fixed-effect layout of a joint fit: the latent stacks all arm betas as a
# contiguous prefix (see .joint_multi_layout), so the generic accessors can
# slice coef/summary to the fixed effects instead of the full latent vector.
.joint_fixed_layout <- function(responses) {
arm_names <- names(responses) %||% paste0("arm", seq_along(responses))
nm <- unlist(lapply(seq_along(responses), function(k) {
X <- responses[[k]]$X
cn <- colnames(X) %||% paste0("beta", seq_len(ncol(X)))
cn[!nzchar(cn)] <- paste0("beta", which(!nzchar(cn)))
paste(arm_names[k], cn, sep = ".")
}), use.names = FALSE)
list(n_fixed = length(nm), names = nm)
}
# Latent-vector layout for the multi-block joint result. Mirrors the
# single-block `.joint_layout()` but generalised over `prepared` blocks:
# per-arm beta, per-arm RE, then each prepared block in order. For
# back-compat with single-block consumers (e.g. tulpaObs cover-hurdle's
# `.joint_field_at_obs_copy_multi`), the first spatial-like block also
# emits `phi_start` / `theta_start` aliases:
#
# * BYM2 -> phi_start, theta_start (length-2 sub-block).
# * ICAR / CAR_proper -> phi_start (length-1).
#
# Non-spatial blocks expose `block_start[b]` so callers can index into
# any block by ordinal position.
.joint_multi_layout <- function(arms, prepared) {
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]
}
B <- length(prepared)
block_start <- integer(B)
block_size <- integer(B)
phi_start <- NULL
theta_start <- NULL
# Per-spatial-block latent offset (0-based) of the structured field, in
# block order: ICAR / CAR_proper at the block start, BYM2 at its phi
# sub-block start. `phi_start` (back-compat) is field_starts[1].
field_starts <- integer(0)
field_block_types <- character(0)
field_n_fields <- integer(0)
for (b in seq_len(B)) {
block_start[b] <- cur
type <- tolower(prepared[[b]]$type %||% "")
# Latent size per block. BYM2 is length-2 (phi ICAR + theta IID)
# on n_spatial_units; all other blocks are length-1 on the
# block's `n_spatial_units` / `n_times` / `n_units`.
n_units <- prepared[[b]]$n_spatial_units %||%
prepared[[b]]$n_times %||%
prepared[[b]]$n_units %||%
prepared[[b]]$n_groups %||%
prepared[[b]]$n_latent %||%
prepared[[b]]$n_mesh %||%
prepared[[b]]$m_total %||% 0L
n_units <- as.integer(n_units)
if (type == "bym2") {
block_size[b] <- 2L * n_units
field_starts <- c(field_starts, cur)
field_block_types <- c(field_block_types, "bym2")
field_n_fields <- c(field_n_fields, 1L)
if (is.null(phi_start)) {
phi_start <- cur
theta_start <- cur + n_units
}
} else if (type == "lf") {
# Latent factor block stores u (n_latent) followed by lambda
# (n_arms). Total size = n_latent + n_arms.
block_size[b] <- n_units + n_arms
} else if (type == "hsgp_mo") {
# Multi-output HSGP stores K * M coefficients in output-major
# order: beta_{k, m} at offset k * M + m. K = n_arms.
block_size[b] <- n_arms * n_units
} else if (type == "mcar") {
# Separable MCAR stores p coupled fields field-major: field a at
# offset a * n_units within the block (size p * n_units).
block_size[b] <- as.integer(prepared[[b]]$n_fields) * n_units
field_starts <- c(field_starts, cur)
field_block_types <- c(field_block_types, "mcar")
field_n_fields <- c(field_n_fields, as.integer(prepared[[b]]$n_fields))
if (is.null(phi_start)) phi_start <- cur
} else if (type == "miid") {
# Multivariate IID stores p coupled per-group coefficient fields
# field-major: field a at offset a * n_groups within the block
# (size p * n_groups). A per-group RE, not a spatial field, so it
# does not register a field_start / phi_start alias.
block_size[b] <- as.integer(prepared[[b]]$n_fields) * n_units
} else {
block_size[b] <- n_units
if (type %in% c("icar", "car_proper")) {
field_starts <- c(field_starts, cur)
field_block_types <- c(field_block_types, type)
field_n_fields <- c(field_n_fields, 1L)
if (is.null(phi_start)) {
phi_start <- cur
}
}
}
cur <- cur + block_size[b]
}
n_x <- cur
out <- list(
n_arms = n_arms,
p = p_arm,
n_re = n_re,
beta_start = beta_start,
re_start = re_start,
block_start = block_start,
block_size = block_size,
n_x = n_x
)
if (!is.null(phi_start)) out$phi_start <- phi_start
if (!is.null(theta_start)) out$theta_start <- theta_start
if (length(field_starts) > 0L) {
out$field_starts <- field_starts
out$field_block_types <- field_block_types
out$field_n_fields <- field_n_fields
}
out
}
# Posterior moments for the multi-block joint grid. Mirrors the single-arm
# multi-block helper (`.nl_posterior_moments_multi`) and additionally
# attaches the weighted-quantile median + 2.5/97.5 CI for every joint-grid
# axis -- the calibrated headline summary for right-skewed scale-like
# hyperparameters (alpha at small n_pos, sigma/range/phi near boundary).
.joint_posterior_moments_multi <- function(res, prepared, axis_offsets,
joint_grid, cp,
int_weights = NULL,
support = .NL_SUPPORT_KINDS,
within = NULL) {
support <- match.arg(support)
within <- .nl_within_cell_mode(within)
w <- res$weights
# Joint moments across every column of joint_grid (including phi
# columns appended for per-arm dispersion overrides).
res$theta_mean <- as.numeric(crossprod(w, joint_grid))
names(res$theta_mean) <- colnames(joint_grid)
res$theta_sd <- sqrt(pmax(0, as.numeric(crossprod(w, joint_grid^2)) -
res$theta_mean^2))
names(res$theta_sd) <- colnames(joint_grid)
B <- length(prepared)
per_block_moments <- vector("list", B)
for (b in seq_len(B)) {
n_axes_b <- axis_offsets[b + 1L] - axis_offsets[b]
if (n_axes_b == 0L) {
# Block contributes no outer-grid axes (e.g. type = "lf"):
# empty moments, empty axis_cols.
per_block_moments[[b]] <- list(
type = tolower(prepared[[b]]$type),
mean = numeric(0),
sd = numeric(0),
axis_cols = integer(0)
)
next
}
cols <- seq.int(axis_offsets[b] + 1L, axis_offsets[b + 1L])
sub <- joint_grid[, cols, drop = FALSE]
block_mean <- as.numeric(crossprod(w, sub))
block_sd <- sqrt(pmax(0, as.numeric(crossprod(w, sub^2)) - block_mean^2))
# Strip the "b<N>." prefix on per-block moment names -- block index
# is implicit in the list position.
bare_names <- sub("^b[0-9]+\\.", "", colnames(sub))
names(block_mean) <- bare_names
names(block_sd) <- bare_names
per_block_moments[[b]] <- list(
type = tolower(prepared[[b]]$type),
mean = block_mean,
sd = block_sd,
axis_cols = cols
)
}
res$block_moments <- per_block_moments
# Each axis's SD from the estimator its own resolution calls for, the rule
# `.nl_posterior_moments()` applies on every other nested path.
res <- .nl_attach_axis_sd(res)
# Weighted-quantile median + 2.5/97.5 CI per axis. After the (sigma,
# alpha) reparameterization, alpha is a column of joint_grid like
# any other axis; this attaches the calibrated summary for every
# axis, not just alpha.
# Supplied whatever the support: a moment rule needs the domain to form its
# interval at all, and a density read needs it to place its outer cell edges
# inside the quantity's own support.
geo <- .joint_axis_geometry(list(theta_grid = joint_grid,
axis_offsets = axis_offsets,
blocks = prepared))
qs <- .nl_axis_quantiles(joint_grid, res$log_marginal,
res$refining_axis, weights = int_weights,
log_quad = if (is.null(int_weights)) res$log_quad,
support = support, domains = geo$domain,
within = within, atoms = geo$atom, sd = TRUE)
res$theta_median <- qs$median
res$theta_ci_lo <- qs$ci_lo
res$theta_ci_hi <- qs$ci_hi
res$within_cell_requested <- within
.nl_attach_interval_provenance(res, qs, joint_grid, geo$domain, geo$atom)
}
Try the tulpa package in your browser
Any scripts or data that you put into this service are public.
tulpa documentation built on Oct. 19, 2026, 5:09 p.m.
R Package Documentation
Browse R Packages
We want your feedback!
Note that we can't provide technical support on individual packages. You should contact the package authors for that.
Defines functions .joint_posterior_moments_multi .joint_multi_layout .joint_fixed_layout .joint_dispatch_multi .nlj_multi_inner_skew_at_theta .joint_attach_pareto_k_multi .joint_multi_make_kernel_fn .joint_multi_call_factory .joint_multi_declared_axes .joint_multi_hyperprior .joint_multi_add_hp .joint_multi_hp_specs .joint_multi_attach_integration .joint_multi_hp_families .joint_multi_measure_specs .joint_check_multi_hyperpriors .joint_multi_hp_cols .joint_block_axis_cols .joint_hp_named_blocks .joint_hp_fn_for_block .joint_parse_hyperprior .joint_is_block_hyperprior .joint_hp_field .joint_multi_phi_per_arm .joint_multi_cross_phi .joint_multi_cpp_grid .multi_block_per_arm_idx .multi_block_svc_weight .joint_block_spec_for_cpp .joint_block_axis_grid .resolve_copy_multi .resolve_one_copy_spec .is_copy_spec_list .normalise_joint_arm_multi
# Multi-block joint dispatch (list-of-blocks prior).
#
# Part of the joint nested-Laplace driver; the public entry point
# tulpa_nested_laplace_joint() lives in nested_laplace_joint.R.
# =============================================================================
# Multi-block joint dispatch (Phase J-B)
# =============================================================================
#
# Activated when `prior` is a list-of-blocks (each element has a `type`
# field). Routes to cpp_nested_laplace_joint_multi, which builds a
# std::vector<LatentBlock> on the joint side. Each block can be a copy
# block; copy (alpha-coupling) is supported on areal spatial (icar / bym2 /
# car_proper), temporal (rw1 / rw2 / ar1), and unstructured (iid) blocks.
# Validate one arm spec for the multi-block path. Same core as
# `.normalise_joint_arm` *except* `spatial_idx` is no longer required at
# the arm level -- per-arm idx vectors live inside each block spec instead,
# and a missing spatial_idx gets a length-N placeholder of zeros so the
# existing JointArm packaging path doesn't complain.
.normalise_joint_arm_multi <- function(a, k) {
.normalise_joint_arm_core(a, k, spatial_idx = "optional")
}
# Detect whether `copy` is a LIST of copy specs (one coupled field each)
# rather than a single spec. A list-of-specs is an unnamed list whose
# elements are themselves lists carrying a `block` field; a single spec is
# itself a list with a `block` (or `arm`) entry at the top level.
.is_copy_spec_list <- function(copy) {
if (!is.list(copy)) return(FALSE)
if (!is.null(copy$block) || !is.null(copy$arm)) return(FALSE)
nm <- names(copy)
has_no_names <- is.null(nm) || all(!nzchar(nm))
has_no_names && length(copy) > 0L &&
all(vapply(copy, function(s) is.list(s) && !is.null(s$block),
logical(1)))
}
# Resolve and validate one copy spec into a (arm_zero, block_zero,
# alpha_grid) triple. The per-spec validation (arm in range, block in
# range, block type spatial, alpha_grid non-negative non-empty) lives here
# so the single-spec and list-of-specs paths share one source of truth.
.resolve_one_copy_spec <- function(spec, responses, prior_list) {
arm_id <- spec$arm
if (is.null(arm_id)) {
stop("`copy$arm` must be a name or 1-based index.", call. = FALSE)
}
if (is.character(arm_id)) {
nm <- names(responses)
if (is.null(nm) || !arm_id %in% nm) {
stop("`copy$arm` = '", arm_id, "' not found in names(responses).",
call. = FALSE)
}
arm_zero <- match(arm_id, nm) - 1L
} else {
arm_zero <- as.integer(arm_id) - 1L
if (arm_zero < 0L || arm_zero >= length(responses)) {
stop("`copy$arm` index out of range.", call. = FALSE)
}
}
block_id <- spec$block
if (is.null(block_id)) {
stop("`copy$block` must be a 1-based index into `prior` for the ",
"multi-block joint driver.", call. = FALSE)
}
block_zero <- as.integer(block_id) - 1L
if (block_zero < 0L || block_zero >= length(prior_list)) {
stop("`copy$block` index (", block_id, ") out of range for length(prior) (",
length(prior_list), ").", call. = FALSE)
}
block_type <- tolower(prior_list[[block_zero + 1L]]$type)
# Copy (alpha-coupling) is supported on any block whose amplitude is a
# single scalar the donor / copy arms scale independently: areal spatial
# (icar / bym2 / car_proper), temporal (rw1 / rw2 / ar1), unstructured
# (iid), and the separable correlated areal field (mcar). For mcar the copy
# scales the WHOLE correlated (intercept, slope) field by one alpha onto the
# copy arm (the cross-arm transfer); the free cross-covariance Sigma stays
# the within-arm covariance among the fields, integrated over the outer grid
#. Blocks whose per-arm scaling lives elsewhere (lf via
# lambda, hsgp_mo via Sigma) or whose precision is precomputed (tgmrf) do
# not take a copy.
copy_types <- c("icar", "bym2", "car_proper", "rw1", "rw2", "ar1", "iid",
"mcar", "miid")
if (!block_type %in% copy_types) {
stop("`copy$block` points at type '", block_type, "'. Copy semantics ",
"are supported on types (",
paste(shQuote(copy_types), collapse = ", "), "). ",
"Blocks with their own per-arm scaling (lf, hsgp_mo) or a ",
"precomputed precision (tgmrf) do not take a copy.", call. = FALSE)
}
# `alpha_grid` states the nodes; `alpha_n` raises the resolution of the
# engine's own axis, which carries 0 (so the "no copy" base model keeps
# posterior mass) plus a log-spaced slab.
alpha_axis <- .nl_copy_alpha_axis(spec[["alpha_grid"]], spec[["alpha_n"]],
what = "`copy`")
if (length(alpha_axis) == 0L) {
stop("`copy$alpha_grid` must have at least one non-negative value.",
call. = FALSE)
}
if (any(alpha_axis < 0)) {
stop("`copy$alpha_grid` values must be non-negative.",
call. = FALSE)
}
list(arm_zero = arm_zero, block_zero = block_zero, alpha_grid = alpha_axis)
}
# Resolve the copy specification for the multi-block joint driver. `copy`
# may be a single spec `list(arm, block, alpha_grid)` (back-compat) OR an
# unnamed list of such specs (N coupled fields, one alpha axis each). Each
# spec must name a distinct copy-supported block. Returns vectors parallel
# over the copy blocks:
# has_copy, copy_arms_zero (int vec), copy_blocks_zero (int vec),
# alpha_grids (list of numeric, one per copy block).
.resolve_copy_multi <- function(copy, responses, prior_list) {
if (is.null(copy)) {
return(list(has_copy = FALSE, copy_arms_zero = integer(0),
copy_blocks_zero = integer(0), alpha_grids = list()))
}
specs <- if (.is_copy_spec_list(copy)) copy else list(copy)
resolved <- lapply(specs, .resolve_one_copy_spec,
responses = responses, prior_list = prior_list)
copy_blocks_zero <- vapply(resolved, function(r) r$block_zero, integer(1))
copy_arms_zero <- vapply(resolved, function(r) r$arm_zero, integer(1))
alpha_grids <- lapply(resolved, function(r) r$alpha_grid)
if (anyDuplicated(copy_blocks_zero)) {
stop("Two copy specs target the same block (1-based index ",
paste(unique(copy_blocks_zero[duplicated(copy_blocks_zero)]) + 1L,
collapse = ", "),
"). Each coupled field must name a distinct block.",
call. = FALSE)
}
list(has_copy = TRUE, copy_arms_zero = copy_arms_zero,
copy_blocks_zero = copy_blocks_zero, alpha_grids = alpha_grids)
}
# Per-block axis grid for the multi-block joint driver. When the block is a
# copy block (any of icar / bym2 / car_proper / rw1 / rw2 / ar1 / iid), the
# parameterisation uses (sigma, alpha[, rho/rho_car]) directly. Non-copy
# blocks use the standard single-arm conventions and reuse the `.NL_REGISTRY`
# defaults.
#
# Returns:
# $grid : matrix [n_block_cells x n_axes_for_block]
# $names : axis names
# $prepared: block spec with defaults filled in (so downstream code can
# read prior$adj_row_ptr etc. without re-checking presence)
.joint_block_axis_grid <- function(p, is_copy, alpha_grid,
block_index) {
type <- tolower(p$type)
if (is_copy && type %in% c("mcar", "miid")) {
# Copied correlated field (mcar: areal; miid: per-group RE) -- the
# block's own axes are the p(p+1)/2 log-Cholesky coordinates of Sigma
# (the within-arm covariance among the fields); the copy appends ONE
# trailing `alpha`
# axis carrying the cross-arm copy amplitude. The kernel reads the
# log-Cholesky axes at axis0 and alpha at axis0 + p(p+1)/2 (the last
# column), so alpha MUST be appended last. Sigma stays a free outer-grid
# quantity (NOT a copy-materialized sigma_occ/sigma_pos pair); alpha is
# applied directly via arm_scale, not multiplied by a sigma.
base <- .nl_block_axis_grid(p)
bg <- base$grid
ag <- as.numeric(alpha_grid)
idx <- expand.grid(b = seq_len(nrow(bg)), a = seq_along(ag),
KEEP.OUT.ATTRS = FALSE)
grid <- cbind(bg[idx$b, , drop = FALSE], alpha = ag[idx$a])
colnames(grid) <- c(colnames(bg), "alpha")
return(list(grid = as.matrix(grid), names = colnames(grid),
prepared = base$prepared))
}
if (is_copy) {
sigma_axis <- as.numeric(p$sigma_grid %||% .nl_grid_axis("field_sd"))
# Copy axes always lead with (sigma, alpha); the kernel materializes
# sigma_donor = sigma and sigma_copy = alpha * sigma in those two
# leading columns (see .joint_multi_cpp_grid), so any trailing
# correlation axis (rho / rho_car) must follow them. A unit-precision
# block (icar / rw1 / rw2 / iid) has just (sigma, alpha); a block with a
# second hyperparameter (bym2 rho, car_proper rho_car, ar1 rho) appends
# it.
if (type %in% c("icar", "rw1", "rw2", "iid")) {
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else if (type == "bym2") {
rho <- p$rho_grid %||% .nl_grid_axis("bym2_rho")
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
rho = as.numeric(rho),
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else if (type == "car_proper") {
rho_car <- p$rho_car_grid %||% .nl_grid_axis("joint_car_rho")
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
rho_car = as.numeric(rho_car),
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else if (type == "ar1") {
rho <- p$rho_grid %||% .nl_grid_axis("ar1_rho")
gr <- expand.grid(sigma = sigma_axis,
alpha = alpha_grid,
rho = as.numeric(rho),
KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)
} else {
stop("Block ", block_index, ": copy semantics are not supported ",
"for type '", type, "'.", call. = FALSE)
}
return(list(grid = as.matrix(gr),
names = colnames(gr),
prepared = p))
}
# Non-copy: reuse the single-arm registry's grid construction. The
# registry returns the standard (tau, [rho]) / (sigma, [rho]) axes.
.nl_block_axis_grid(p)
}
# Convert one R-side block to the C++ `cpp_nested_laplace_joint_multi`
# format. Per-arm idx vectors live in the BLOCK spec (not the arm spec);
# the user supplies `spatial_idx` / `temporal_idx` / `obs_idx` as a list
# of length n_arms.
#
# `arms` is the list of normalised arm specs produced by
# `.normalise_joint_arm_multi()`. It is consumed by block types that need
# to read covariate columns at fit-time (HSGP with `svc_column`). Pass
# NULL only from contexts where no such block can appear.
.joint_block_spec_for_cpp <- function(p, n_arms, block_index, arms = NULL) {
type <- tolower(p$type)
projected <- .nl_block_is_projected(p)
.nl_check_block_fields(
p, c("axis", if (projected) "projected" else "joint"), block_index)
arm_n_obs <- if (is.null(arms)) NULL
else vapply(arms, function(a) length(a$y), integer(1))
if (type %in% .NL_PROJECTABLE_AREAL) {
spatial_idx <- if (!projected)
.multi_block_per_arm_idx(p$spatial_idx, n_arms, block_index,
"spatial_idx", arm_n_obs,
p$n_spatial_units)
out <- list(
type = type,
spatial_idx = spatial_idx,
n_spatial_units = as.integer(p$n_spatial_units),
adj_row_ptr = as.integer(p$adj_row_ptr),
adj_col_idx = as.integer(p$adj_col_idx),
n_neighbors = as.integer(p$n_neighbors)
)
if (projected) {
out$projector <- .nl_block_projector(p, n_arms, block_index,
arm_n_obs)
}
if (type == "bym2") {
out$scale_factor <- as.numeric(p$scale_factor %||% 1.0)
if (!is.null(p$node_prec)) out$node_prec <- as.numeric(p$node_prec)
}
# Optional per-arm per-row design weight (areal SVC). When set, the
# field's contribution to arm k row i is row-scaled by svc_weight[[k]][i]:
# eta_i += svc_weight[[k]][i] * amplitude * z[cell_i].
# Parallel to `spatial_idx`: a list of length n_arms, each a numeric
# vector matching that arm's obs count and the per-arm spatial_idx
# length. Composes multiplicatively with copy `arm_scale`.
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, spatial_idx, n_arms, block_index)
# Aliased fixed-effect column of a weighted (areal SVC) field: its
# constant folds into this coefficient (0 = intercept, or the
# covariate's column) rather than the intercept; -1 = no counterpart
# (the level is left in the field). Absent => no fold.
if (!is.null(p$svc_beta_offset)) {
out$svc_beta_offset <- as.integer(p$svc_beta_offset)
}
}
# Replicated CAR (`by =`): number of disjoint equal-size components in the
# block-diagonal graph (the intrinsic ICAR kernel pins each per component
# and uses (n - L) in its rank-deficiency normalizer). Absent = connected.
if (!is.null(p$n_components)) {
out$n_components <- as.integer(p$n_components)
}
out
} else if (type == "mcar") {
# Separable multivariate CAR: p coupled areal fields sharing Sigma (x)
# Q^-1. One block over p * n_spatial_units latent. spatial_idx is the
# per-arm 1-based cell index (shared by all fields); field_weight is the
# per-field design column (the intercept's is all-ones, a covariate
# column is the per-row value) -- outer length p, inner length n_arms.
spatial_idx <- .multi_block_per_arm_idx(p$spatial_idx, n_arms,
block_index, "spatial_idx",
arm_n_obs, p$n_spatial_units)
n_fields <- as.integer(p$n_fields)
if (length(p$field_weight) != n_fields) {
stop("Block ", block_index, " (type 'mcar'): `field_weight` must be ",
"a list of length n_fields (", n_fields, ").", call. = FALSE)
}
field_weight <- lapply(seq_len(n_fields), function(a) {
.multi_block_svc_weight(p$field_weight[[a]], spatial_idx, n_arms,
block_index)
})
out <- list(
type = "mcar",
spatial_idx = spatial_idx,
n_spatial_units = as.integer(p$n_spatial_units),
n_fields = n_fields,
adj_row_ptr = as.integer(p$adj_row_ptr),
adj_col_idx = as.integer(p$adj_col_idx),
n_neighbors = as.integer(p$n_neighbors),
field_weight = field_weight
)
# Per-field aliased fixed-effect column (length n_fields): field a's
# global constant folds into field_beta_offset[a] (0 = intercept for the
# all-ones field, the covariate's column otherwise, -1 = no counterpart).
# Absent => no folds (the augmentation identifies every level in place).
if (!is.null(p$field_beta_offset)) {
if (length(p$field_beta_offset) != n_fields) {
stop("Block ", block_index, " (type 'mcar'): `field_beta_offset` ",
"must have length n_fields (", n_fields, ").", call. = FALSE)
}
out$field_beta_offset <- as.integer(p$field_beta_offset)
}
# Replicated MCAR (`by =`): L equal-size components per field.
if (!is.null(p$n_components)) {
out$n_components <- as.integer(p$n_components)
}
out
} else if (type %in% c("rw1", "rw2", "ar1")) {
temporal_idx <- .multi_block_per_arm_idx(p$temporal_idx, n_arms,
block_index, "temporal_idx",
arm_n_obs, p$n_times)
out <- list(
type = type,
temporal_idx = temporal_idx,
n_times = as.integer(p$n_times)
)
if (type %in% c("rw1", "rw2")) out$cyclic <- isTRUE(p$cyclic)
# Optional per-arm per-row design weight (temporal SVC). Parallel to the
# areal branch: the field's contribution to arm k row i is row-scaled by
# svc_weight[[k]][i] (eta_i += svc_weight[[k]][i] * amplitude *
# f[time_i]), used for a temporally varying slope on a covariate column.
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, temporal_idx, n_arms, block_index)
# Aliased fixed-effect column of a weighted (temporal SVC / TVC)
# field, as in the areal branch above; absent => no fold.
if (!is.null(p$svc_beta_offset)) {
out$svc_beta_offset <- as.integer(p$svc_beta_offset)
}
}
out
} else if (type == "iid") {
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, p$n_units)
out <- list(
type = "iid",
obs_idx = obs_idx,
n_units = as.integer(p$n_units)
)
# Optional per-arm per-row design weight (random slope on a covariate
# column). When set, the field's contribution to arm k row i is
# row-scaled by svc_weight[[k]][i]:
# eta_i += svc_weight[[k]][i] * sigma * u[obs_idx_i].
# Parallel to the areal / temporal SVC branches: a list of length n_arms,
# each a numeric vector matching that arm's obs count (the per-arm obs_idx
# length). Composes multiplicatively with copy `arm_scale`. An
# uncorrelated slope (x || g) / (0 + x | g) is one weighted iid block per
# coefficient, each with its own sigma axis.
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, obs_idx, n_arms, block_index)
}
out
} else if (type == "miid") {
# Multivariate IID: p coupled per-group coefficient
# fields sharing a free Sigma (the Q = I sibling of mcar). One block over
# p * n_groups latent. `obs_idx` is the per-arm 1-based grouping index
# (the mcar `spatial_idx` role); `field_weight` is the per-coefficient
# design column (intercept all-ones, slope = covariate) -- outer length
# n_fields, inner length n_arms. Expresses a correlated random slope
# (1 + x | g): n_fields = 1 + n_slopes.
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, p$n_groups)
n_fields <- as.integer(p$n_fields)
if (length(p$field_weight) != n_fields) {
stop("Block ", block_index, " (type 'miid'): `field_weight` must be ",
"a list of length n_fields (", n_fields, ").", call. = FALSE)
}
field_weight <- lapply(seq_len(n_fields), function(a) {
.multi_block_svc_weight(p$field_weight[[a]], obs_idx, n_arms,
block_index)
})
list(
type = "miid",
obs_idx = obs_idx,
n_groups = as.integer(p$n_groups),
n_fields = n_fields,
field_weight = field_weight
)
} else if (type == "tgmrf") {
# User-supplied GMRF: per-grid CSC Q + logdet + log p(theta).
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, p$n_latent)
list(
type = "tgmrf",
n_latent = as.integer(p$n_latent),
obs_idx = obs_idx,
Q_csc_p_per_grid = lapply(p$Q_csc_p_per_grid, as.integer),
Q_csc_i_per_grid = lapply(p$Q_csc_i_per_grid, as.integer),
Q_csc_x_per_grid = lapply(p$Q_csc_x_per_grid, as.numeric),
logdet_Q_per_grid = as.numeric(p$logdet_Q_per_grid),
log_prior_theta_per_grid = as.numeric(p$log_prior_theta_per_grid)
)
} else if (type == "nngp") {
spatial_idx <- if (!projected)
.multi_block_per_arm_idx(p$spatial_idx, n_arms, block_index,
"spatial_idx", arm_n_obs, p$n_spatial)
out <- c(list(type = "nngp", spatial_idx = spatial_idx),
.nl_nngp_field_args(p))
if (projected) {
out$projector <- .nl_block_projector(p, n_arms, block_index,
arm_n_obs)
}
if (!is.null(p$svc_weight)) {
out$svc_weight <- .multi_block_svc_weight(
p$svc_weight, spatial_idx, n_arms, block_index)
}
out
} else if (type == "hsgp") {
# HSGP block: shared eigenvalues + per-arm Phi basis matrices.
# Axes are (sigma2, lengthscale).
#
# Optional `svc_column`: 1-based index into each arm's `X`. When
# set, the block becomes a spatially-varying coefficient:
# eta_i += X[i, svc_column] * f(s_i)
# Implemented by row-scaling the per-arm basis at fit-time
# phi_scaled[k][i, m] = phi[k][i, m] * X[k][i, svc_column]
# and passing the scaled phi to the existing DENSE_BASIS factory.
# Multi-scale composition (1.6c) falls out by declaring multiple
# type='hsgp' blocks in the same prior_list.
if (!is.list(p$phi) || length(p$phi) != n_arms) {
stop("Block ", block_index, " (type 'hsgp'): `phi` must be a ",
"list of length n_arms (", n_arms, ").", call. = FALSE)
}
phi <- lapply(p$phi, function(M) {
M <- as.matrix(M); storage.mode(M) <- "double"; M
})
if (!is.null(p$svc_column)) {
if (is.null(arms)) {
stop("Block ", block_index, " (type 'hsgp'): `svc_column` ",
"requires the joint multi-block dispatcher path ",
"(arms not available here).", call. = FALSE)
}
svc <- as.integer(p$svc_column)
if (length(svc) != 1L || is.na(svc) || svc < 1L) {
stop("Block ", block_index, " (type 'hsgp'): `svc_column` ",
"must be a single positive integer (1-based column ",
"index into each arm's X).", call. = FALSE)
}
for (k in seq_len(n_arms)) {
Xk <- arms[[k]]$X
if (svc > ncol(Xk)) {
stop("Block ", block_index, " (type 'hsgp'): ",
"`svc_column` = ", svc, " exceeds ncol(arm ", k,
"$X) = ", ncol(Xk), ".", call. = FALSE)
}
if (nrow(phi[[k]]) != nrow(Xk)) {
stop("Block ", block_index, " (type 'hsgp'): ",
"nrow(phi[[", k, "]]) (", nrow(phi[[k]]),
") must equal nrow(arm ", k, "$X) (", nrow(Xk),
") when svc_column is set.", call. = FALSE)
}
phi[[k]] <- phi[[k]] * as.numeric(Xk[, svc])
}
}
list(
type = "hsgp",
m_total = as.integer(p$m_total),
phi = phi,
n_obs_per_arm = as.integer(p$n_obs_per_arm),
eigenvalues = as.numeric(p$eigenvalues)
)
} else if (type == "hsgp_mo") {
# Multi-output (co-regionalization) HSGP block. K = n_arms
# correlated latent fields share basis + eigenvalues; coefficients are
# cross-output-correlated via Sigma. First ship restricts K == 2 with
# axes (sigma_1, sigma_2, rho, ell), all raw.
if (n_arms != 2L) {
stop("Block ", block_index, " (type 'hsgp_mo'): ",
"requires n_arms == 2 (got ", n_arms, "). Multi-output HSGP ",
"with K > 2 needs an LKJ-cholesky correlation ",
"parameterization on the outer grid, which is not supported.",
call. = FALSE)
}
if (!is.list(p$phi) || length(p$phi) != n_arms) {
stop("Block ", block_index, " (type 'hsgp_mo'): `phi` must be a ",
"list of length n_arms (", n_arms, ").", call. = FALSE)
}
phi <- lapply(p$phi, function(M) {
M <- as.matrix(M); storage.mode(M) <- "double"; M
})
list(
type = "hsgp_mo",
m_total = as.integer(p$m_total),
phi = phi,
n_obs_per_arm = as.integer(p$n_obs_per_arm),
eigenvalues = as.numeric(p$eigenvalues)
)
} else if (type == "spde") {
# SPDE block: per-arm sparse barycentric projector A and shared FEM
# mesh (C0_diag, G1). Axes are (range, sigma). Optional rational
# coefficients for fractional nu.
per_arm <- function(field) {
v <- p[[field]]
if (!is.list(v) || length(v) != n_arms) {
stop("Block ", block_index, " (type 'spde'): `", field,
"` must be a list of length n_arms (", n_arms, ").",
call. = FALSE)
}
v
}
out <- list(
type = "spde",
n_mesh = as.integer(p$n_mesh),
A_x = lapply(per_arm("A_x"), as.numeric),
A_i = lapply(per_arm("A_i"), as.integer),
A_p = lapply(per_arm("A_p"), as.integer),
n_obs_per_arm = as.integer(p$n_obs_per_arm),
C0_diag = as.numeric(p$C0_diag),
G1_x = as.numeric(p$G1_x),
G1_i = as.integer(p$G1_i),
G1_p = as.integer(p$G1_p),
nu = as.numeric(p$nu)
)
if (!is.null(p$rational_poles))
out$rational_poles <- as.numeric(p$rational_poles)
if (!is.null(p$rational_weights))
out$rational_weights <- as.numeric(p$rational_weights)
out
} else if (type == "lf") {
# Latent factor block. Per-arm `obs_idx` maps each
# obs to a slot in the factor field `u` (length n_latent); the
# per-arm loadings `lambda` (length n_arms) are part of the joint
# latent vector and co-optimized by the inner Newton. No outer-
# grid axes -- identifiability handled by tight Gaussian anchors
# on (u_1, lambda_1) inside the C++ factory.
#
# The unit count is checked BEFORE the indices are checked against it:
# an `n_latent` of 1 makes every real index out of range, so validating
# the indices first reports the block's shape error as an index error.
# The C++ factory keeps its own copy of this check for a consumer that
# builds the block through `LinkingTo` rather than this door.
n_latent <- as.integer(p$n_latent)
if (length(n_latent) != 1L || is.na(n_latent) || n_latent < 2L) {
stop("Block ", block_index, " (type 'lf'): `n_latent` must be at ",
"least 2 (got ", format(p$n_latent), "). The factor field ",
"needs a second slot for the anchor on (u_1, lambda_1) to ",
"leave anything free.", call. = FALSE)
}
obs_idx <- .multi_block_per_arm_idx(p$obs_idx, n_arms,
block_index, "obs_idx",
arm_n_obs, n_latent)
out <- list(
type = "lf",
n_latent = n_latent,
obs_idx = obs_idx
)
if (!is.null(p$sigma_u)) out$sigma_u <- as.numeric(p$sigma_u)
if (!is.null(p$sigma_lambda)) out$sigma_lambda <- as.numeric(p$sigma_lambda)
if (!is.null(p$anchor_eps)) out$anchor_eps <- as.numeric(p$anchor_eps)
out
} else {
stop("Block type '", type, "' is not supported in multi-block joint priors.",
call. = FALSE)
}
}
# Coerce + validate the optional `svc_weight` for an areal block into a
# length-n_arms list of numeric vectors. Each entry must match that arm's
# obs count, which is the length of that arm's resolved `spatial_idx`.
# Accepts either a list (already per-arm) or a single numeric vector
# (replicated across arms when every arm shares the same per-row weight).
.multi_block_svc_weight <- function(svc_weight, spatial_idx, n_arms,
block_index) {
if (is.list(svc_weight)) {
if (length(svc_weight) != n_arms) {
stop("Block ", block_index, ": `svc_weight` is a list of ",
length(svc_weight), " entries but n_arms = ", n_arms, ".",
call. = FALSE)
}
w_list <- lapply(svc_weight, as.numeric)
} else {
w_list <- rep(list(as.numeric(svc_weight)), n_arms)
}
for (k in seq_len(n_arms)) {
n_k <- length(spatial_idx[[k]])
if (length(w_list[[k]]) != n_k) {
stop("Block ", block_index, ": `svc_weight[[", k, "]]` has length ",
length(w_list[[k]]), " but arm ", k, " has ", n_k,
" observations (matching its `spatial_idx`).", call. = FALSE)
}
if (any(!is.finite(w_list[[k]]))) {
stop("Block ", block_index, ": `svc_weight[[", k, "]]` contains ",
"non-finite values.", call. = FALSE)
}
}
w_list
}
# Coerce a per-arm idx field into a length-n_arms list of integer vectors.
# Accepts either a list (already per-arm) or a single vector (replicated
# across arms; useful when every arm shares the same indexing -- e.g. years
# 1..T mapped one-to-one).
# `arm_n_obs`, when supplied, is each arm's observation count
# (length(arms[[k]]$y)): the C++ side (make_per_arm_idx_fn,
# src/nested_laplace_joint_multi.cpp) reads `cache[k_arm][i]` for every row i
# of that arm with NO bounds check, so a per-arm entry shorter than that arm's
# row count is a read out of bounds -- an access violation, not a clean error
# -- the moment a grid cell scatters that arm's rows into this block. There is
# no "this block skips arm k" shorthand via a short/empty vector: a block
# that should not affect an arm still needs a full-length index vector for
# it; exclude the arm's CONTRIBUTION instead (e.g. a field_coef of 0 on that
# arm), which the engine already tests via `d_eff_cache == 0`.
#
# `n_units`, when supplied, is the block's own latent unit count
# (n_spatial_units / n_times / n_units / n_groups / n_latent). Every index must
# land in 1..n_units. The kernel treats an out-of-range index as "this row does
# not see this block" and moves on at every one of its six consumers, so a
# 0-based vector, an off-by-one, or an NA does not fail: the affected rows drop
# out of eta, out of the gradient and out of the Hessian, the inner Newton
# converges on the reduced problem, `converged` is TRUE, and the log-marginal
# surface is the one for a different model. The range is the only place that
# distinguishes a deliberately absent row from a wrong one, so it is checked
# here, once, where the length already is.
.multi_block_per_arm_idx <- function(idx, n_arms, block_index, field_name,
arm_n_obs = NULL, n_units = NULL) {
out <- if (is.list(idx)) {
if (length(idx) != n_arms) {
stop("Block ", block_index, ": `", field_name, "` is a list of ",
length(idx), " entries but n_arms = ", n_arms, ".",
call. = FALSE)
}
lapply(idx, as.integer)
} else {
rep(list(as.integer(idx)), n_arms)
}
if (!is.null(arm_n_obs)) {
for (k in seq_len(n_arms)) {
if (length(out[[k]]) != arm_n_obs[k]) {
stop("Block ", block_index, ": `", field_name, "[[", k, "]]` has ",
length(out[[k]]), " entries but arm ", k, " has ",
arm_n_obs[k], " observations -- every arm needs a ",
"matching-length index vector for every block (there is ",
"no arm-exclusion shorthand). To exclude arm ", k,
" from this block, zero its CONTRIBUTION (e.g. a ",
"`field_coef` of 0 on that arm) rather than shortening ",
"`", field_name, "`.", call. = FALSE)
}
}
}
if (!is.null(n_units)) {
n_units <- as.integer(n_units)
if (is.na(n_units) || n_units < 1L) {
stop("Block ", block_index, ": the latent unit count backing `",
field_name, "` must be a positive integer, got ", n_units, ".",
call. = FALSE)
}
for (k in seq_len(n_arms)) {
v <- out[[k]]
if (length(v) == 0L) next
# The kernel's contract for a block index is that `l > 0 && l <=
# size` contributes and anything else is skipped (LatentBlock::idx,
# and every scatter that reads it). A non-positive entry is
# therefore the documented "this observation does not reach this
# block" sentinel -- what an arm-exclusive block spells on the arms
# it does not reach -- and is accepted. An entry ABOVE n_units names
# a latent unit that does not exist: the kernel drops it just as
# silently, and there is no reading under which that was meant.
bad <- is.na(v) | v > n_units
if (any(bad)) {
i <- which(bad)[1L]
stop("Block ", block_index, ": `", field_name, "[[", k, "]][",
i, "]` is ", v[i], ", above the block's unit count ",
n_units, " (", sum(bad), " of ", length(v),
" entries are). Every positive index must name one of the ",
"block's latent units; an out-of-range one is silently ",
"dropped by the kernel, which fits a different model and ",
"reports convergence. Use 0 (or any non-positive value) ",
"for an observation that does not reach this block.",
call. = FALSE)
}
# A COMPLETE 0-based enumeration is the one pattern likelier to be
# an off-by-one than a sentinel: every unit but the last named, and
# the unnamed slot filled by the value that also means "skip". A
# sentinel pattern reaching exactly 1..n_units-1 and never n_units
# is possible, so this warns rather than rejecting.
if (n_units > 1L && any(v == 0L) &&
setequal(unique(v), seq.int(0L, n_units - 1L))) {
warning("Block ", block_index, ": `", field_name, "[[", k,
"]]` takes every value in 0..", n_units - 1L,
" and never ", n_units, ", the signature of a 0-based ",
"index vector. Indices are 1-based and 0 means \"this ",
"observation does not reach this block\", so unit ",
n_units, " would go unreached and every other ",
"observation would land one unit low.", call. = FALSE)
}
}
}
out
}
# Materialize the C++-facing theta_grid from a user-facing multi-block
# `joint_grid` (the latent-block axis columns; any trailing phi_<arm>
# dispersion columns are excluded here -- they ride on `phi_grid_per_arm`).
# The kernel reads `sigma_occ` / `sigma_pos` on a copy block; the R-side grid
# lives in (sigma, alpha), so this materializes sigma_occ = sigma and
# sigma_pos = alpha * sigma per copy block. Non-copy block columns pass
# through unchanged. Single source of truth for the main dispatch and the
# Pareto-k re-evaluation.
.joint_multi_cpp_grid <- function(joint_grid, axis_offsets, B, cp) {
cpp_grid <- joint_grid[, seq_len(axis_offsets[B + 1L]), drop = FALSE]
if (!cp$has_copy) return(cpp_grid)
new_names <- colnames(cpp_grid)
for (b_copy in (cp$copy_blocks_zero + 1L)) {
cols_b <- (axis_offsets[b_copy] + 1L):axis_offsets[b_copy + 1L]
bare_b <- sub("^b[0-9]+\\.", "", colnames(cpp_grid)[cols_b])
i_sigma <- match("sigma", bare_b)
i_alpha <- match("alpha", bare_b)
# A copied correlated field (mcar) has no scalar
# `sigma` axis -- its amplitude lives in the free Sigma (the log-Cholesky
# axes), and the copy applies alpha directly via arm_scale (NOT
# sigma_occ = sigma / sigma_pos = alpha*sigma). The alpha column passes
# through to the kernel unchanged; the log-Cholesky axes pass through too.
if (is.na(i_sigma) && !is.na(i_alpha)) next
if (is.na(i_sigma) || is.na(i_alpha)) {
stop(".joint_multi_cpp_grid: copy block missing 'sigma' or 'alpha' axis.",
call. = FALSE)
}
sigma_col <- as.numeric(cpp_grid[, cols_b[i_sigma]])
alpha_col <- as.numeric(cpp_grid[, cols_b[i_alpha]])
cpp_grid[, cols_b[i_sigma]] <- sigma_col # -> sigma_occ
cpp_grid[, cols_b[i_alpha]] <- alpha_col * sigma_col # -> sigma_pos
new_names[cols_b[i_sigma]] <- paste0("b", b_copy, ".sigma_occ")
new_names[cols_b[i_alpha]] <- paste0("b", b_copy, ".sigma_pos")
}
colnames(cpp_grid) <- new_names
cpp_grid
}
# Cross the active per-arm dispersion axes of `phi_axes` (the
# `.normalise_phi_grid()` list) onto a latent outer grid as a Cartesian product:
# every latent row repeated once per dispersion cell, latent rows varying
# fastest, the dispersion columns appended as `phi_<arm>`. Returns the crossed
# grid and `latent_row`, the latent row each crossed cell repeats. The cell
# layout depends only on the latent row count and each axis's node count, so
# two responses whose axes share node counts share it, whatever the nodes are.
.joint_multi_cross_phi <- function(joint_grid, phi_axes, axis_names) {
active <- phi_axes[vapply(phi_axes, length, integer(1)) > 0L]
phi_extra <- do.call(expand.grid,
c(active, list(KEEP.OUT.ATTRS = FALSE,
stringsAsFactors = FALSE)))
joint_idx <- expand.grid(joint = seq_len(nrow(joint_grid)),
phi = seq_len(nrow(phi_extra)),
KEEP.OUT.ATTRS = FALSE)
grid <- cbind(joint_grid[joint_idx$joint, , drop = FALSE],
as.matrix(phi_extra[joint_idx$phi, , drop = FALSE]))
colnames(grid) <- c(axis_names, paste0("phi_", names(active)))
list(grid = grid, latent_row = joint_idx$joint)
}
# Extract the per-arm dispersion override list from a `joint_grid` carrying
# `phi_<arm>` columns. Returns NULL when no such column exists (the kernel
# then uses each arm's parse-time scalar phi); otherwise a length-n_arms list
# with entry k the column `phi_<arm_names[k]>` or NULL.
.joint_multi_phi_per_arm <- function(joint_grid, arm_names) {
cn <- colnames(joint_grid)
if (is.null(cn) || !any(startsWith(cn, "phi_"))) return(NULL)
out <- vector("list", length(arm_names))
for (k in seq_along(arm_names)) {
col <- paste0("phi_", arm_names[k])
if (col %in% cn) out[[k]] <- as.numeric(joint_grid[, col])
}
out
}
# --- regularizing hyperpriors on a multi-block grid ------------------------
#
# `prior_sigma` and `prior_alpha` name a ROLE -- the donor field amplitude and
# the copy coefficient -- and a multi-block grid can carry that role on several
# blocks at once: an occurrence field copied onto two arms carries `b1.alpha`
# and `b2.alpha`. Two shapes state which axes one density reaches.
#
# * ONE spec, `list(<family>, <params>)`: the density applies to EVERY block
# axis of that role. That is the reading `prior_phi` already takes across
# the per-arm dispersion axes, and the one a single-block grid's single
# axis cannot distinguish. Each axis carries its own copy of the density
# and each axis's atom / continuum split is normalised against that axis's
# own cell widths -- `.hyper_atom_fold_scale()` reads one spec's own
# `log_prior` against its own continuum -- so N axes take N independent
# applications and no normalizer is shared between blocks.
# * A LIST of per-block entries, `list(list(block = , prior = ), ...)`, keyed
# by 1-based block index the way a multi-block `copy` spec keys its blocks.
# A block the list does not name carries no density on that axis.
#
# Parsing returns NULL (flat), one density function (the first shape), or a
# `tulpa_block_hyperprior`: density functions named by block index.
# One named element of a list, NULL where the list carries no such name.
.joint_hp_field <- function(x, name) {
i <- match(name, names(x))
if (is.na(i)) NULL else x[[i]]
}
# Is `spec` the per-block shape? Every element of the scalar shape is a family
# name or a parameter vector; every element of this one is a list naming a
# block.
.joint_is_block_hyperprior <- function(spec) {
if (!is.list(spec) || length(spec) == 0L) return(FALSE)
all(vapply(spec, function(e) is.list(e) && "block" %in% names(e),
logical(1)))
}
# Parse a `prior_sigma` / `prior_alpha` argument. `multi_block` says whether
# the fit's `prior` is a list of blocks: the per-block shape names blocks of
# such a list, so on the single-block path it is refused rather than read as a
# malformed family spec.
.joint_parse_hyperprior <- function(spec, axis_label, multi_block = FALSE) {
if (is.null(spec)) return(NULL)
if (!.joint_is_block_hyperprior(spec)) {
return(.joint_parse_sigma_prior(spec, axis_label))
}
if (!isTRUE(multi_block)) {
stop("`", axis_label, "` is a per-block list, which names blocks of a ",
"list-of-blocks `prior`; this fit takes the single-block joint ",
"path, where the axis is unique. Pass the spec itself, e.g. ",
"list(\"pc.prec\", c(U, alpha)).", call. = FALSE)
}
out <- list()
for (i in seq_along(spec)) {
e <- spec[[i]]
b <- .joint_hp_field(e, "block")
if (length(b) != 1L || !is.numeric(b) || !is.finite(b) ||
b != round(b) || b < 1) {
stop("`", axis_label, "[[", i, "]]$block` must be a single 1-based ",
"block index into `prior`.", call. = FALSE)
}
key <- as.character(as.integer(b))
if (key %in% names(out)) {
stop("`", axis_label, "` names block ", key, " twice. Each block ",
"takes at most one hyperprior on a given axis.", call. = FALSE)
}
p <- .joint_hp_field(e, "prior")
if (is.null(p)) {
stop("`", axis_label, "[[", i, "]]` names block ", key, " but ",
"carries no `prior` field. Each entry is ",
"list(block = <index>, prior = list(<family>, <params>)).",
call. = FALSE)
}
out[[key]] <- .joint_parse_sigma_prior(
p, sprintf("%s[[%d]]$prior", axis_label, i))
}
structure(out, class = "tulpa_block_hyperprior")
}
# The density one hyperprior applies to block `b`: the stated density wherever
# the caller stated one, that block's own entry where the caller keyed the
# prior per block, NULL where neither reaches this block.
.joint_hp_fn_for_block <- function(fn, b) {
if (is.null(fn)) return(NULL)
if (is.function(fn)) return(fn)
.joint_hp_field(fn, as.character(b))
}
# The block indices a per-block hyperprior names. A stated single density names
# none -- it reaches every block carrying the axis -- so this is `integer(0)`
# for it as well as for an absent prior.
.joint_hp_named_blocks <- function(fn) {
if (is.null(fn) || is.function(fn)) return(integer(0))
as.integer(names(fn))
}
# The grid column carrying one axis role in each block that has it, named by
# 1-based block index. A block with no axes at all (a latent-factor block
# carries a 1 x 0 grid) contributes none: its column span is empty, which
# `seq.int(length.out =)` reads as empty where a `from:to` range would run
# backwards into the neighbouring block's columns.
.joint_block_axis_cols <- function(axis_names, axis_offsets, B, role) {
out <- integer(0)
for (b_idx in seq_len(B)) {
n_b <- axis_offsets[b_idx + 1L] - axis_offsets[b_idx]
if (n_b <= 0L) next
cols_b <- seq.int(axis_offsets[b_idx] + 1L, length.out = n_b)
i <- match(role, sub("^b[0-9]+\\.", "", axis_names[cols_b]))
if (is.na(i)) next
out[as.character(b_idx)] <- cols_b[i]
}
out
}
# Which columns of a multi-block `joint_grid` each hyperprior applies to, and
# with which density. One entry per applied (column, role): `col` indexes the
# grid, `role` is `"sigma"` / `"alpha"` / `"phi"`, `block` is the 1-based block
# index (NA for a per-arm phi axis) and `fn` the density folded onto it.
# Single source of the mapping, so the fold below and the axis quadrature
# (`.joint_multi_hp_specs()`) name the same columns and the same densities.
.joint_multi_hp_cols <- function(joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi = NULL) {
cn <- colnames(joint_grid)
out <- list()
fns <- list(sigma = fn_sigma, alpha = fn_alpha)
for (role in c("sigma", "alpha")) {
if (is.null(fns[[role]])) next
cols <- .joint_block_axis_cols(cn, axis_offsets, B, role)
for (k in seq_along(cols)) {
b_idx <- as.integer(names(cols)[k])
fn <- .joint_hp_fn_for_block(fns[[role]], b_idx)
if (is.null(fn)) next
out[[length(out) + 1L]] <- list(col = as.integer(cols[[k]]),
role = role, block = b_idx, fn = fn)
}
}
# phi_<arm> dispersion axes are not block-prefixed; carry them straight
# through so a single `fn_phi` re-weights each on the joint grid.
if (!is.null(fn_phi)) {
for (j in grep("^phi_", cn)) {
out[[length(out) + 1L]] <- list(col = as.integer(j), role = "phi",
block = NA_integer_, fn = fn_phi)
}
}
out
}
# Which axes each regularizing hyperprior reaches, and refusal of a per-block
# entry naming a block that carries no axis of that role. Read off the axis
# layout, so a mis-keyed prior is refused before the first inner solve rather
# than applied to nothing. Returns NULL when no hyperprior is active, else one
# record per role: the scope the caller stated it at, the blocks it reaches and
# the grid columns those are.
.joint_check_multi_hyperpriors <- function(axis_names, axis_offsets, B,
fn_sigma, fn_alpha) {
roles <- list(prior_sigma = "sigma", prior_alpha = "alpha")
fns <- list(prior_sigma = fn_sigma, prior_alpha = fn_alpha)
out <- list()
for (label in names(roles)) {
fn <- fns[[label]]
if (is.null(fn)) next
role <- roles[[label]]
cols <- .joint_block_axis_cols(axis_names, axis_offsets, B, role)
carried <- as.integer(names(cols))
unnamed <- setdiff(.joint_hp_named_blocks(fn), carried)
if (length(unnamed)) {
stop("`", label, "` names block",
if (length(unnamed) > 1L) "s " else " ",
paste(unnamed, collapse = ", "),
", which carr", if (length(unnamed) > 1L) "y" else "ies",
" no `", role, "` axis. ",
if (length(carried))
paste0("The blocks carrying one are ",
paste(carried, collapse = ", "),
" (grid column",
if (length(carried) > 1L) "s " else " ",
paste(axis_names[cols], collapse = ", "), ").")
else paste0("No block of this fit carries a `", role,
"` axis."),
call. = FALSE)
}
applied <- if (is.function(fn)) carried
else intersect(carried, .joint_hp_named_blocks(fn))
out[[role]] <- list(
scope = if (is.function(fn)) "every_block" else "per_block",
blocks = as.integer(applied),
axes = as.character(axis_names[cols[as.character(applied)]]))
}
if (length(out) == 0L) NULL else out
}
# The axis specs a multi-block joint grid's cell measure is built from: the
# specs read off `grid`'s columns with the folded axes marked, the copy scale's
# declared atom mass, and the caller's hyperprior roles
# (`.joint_multi_hp_specs()`), with each free-covariance block measured under
# the design `logchol` declares (`.joint_multi_declared_axes()`). The one
# construction behind the weights and the screen that ranks the cells before any
# is solved.
.joint_multi_measure_specs <- function(grid, folded_axes, joint_grid,
axis_offsets, B, fn_sigma, fn_alpha,
fn_phi, copy_slab, copy_atom_mass,
logchol, axis_refine = NULL) {
specs <- .joint_axis_specs_from_grid(grid, copy_slab = copy_slab,
folded_axes = folded_axes,
logchol = logchol,
axis_refine = axis_refine)
specs <- lapply(specs, function(sp) {
if (!is.null(sp$atom_mass)) sp$atom_mass <- copy_atom_mass
sp
})
.joint_multi_hp_specs(specs, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi)
}
# The hyperprior family of each arm, keyed by arm name, as the default
# hyperprior reads it.
.joint_multi_hp_families <- function(arms, arm_names) {
stats::setNames(
vapply(arms, function(a) as.character(a$family %||% ""), character(1)),
arm_names)
}
# The outer integration of a multi-block joint result whose `log_marginal`
# already carries the folded hyperprior: the hyperprior record over
# `joint_grid`, the cell measure `log_quad` and the axis supports over
# `res$theta_grid` from the specs that record marks, and the integration
# weights, with `dnode` the design weights of a CCD design (NULL on a tensor
# grid). Returns the result and the specs its measure was built from.
#
# `refining` is the per-cell tag the refinement passes leave. The specs are read
# off the grid's BASE cells: an axis's declared node set is what fixes a prior
# read off it, and a node a refinement pass appended was not declared.
.joint_multi_attach_integration <- function(res, joint_grid, axis_offsets, B,
prepared, hp_families,
fn_sigma, fn_alpha, fn_phi,
copy_slab, copy_atom_mass,
dnode = NULL,
hyperprior = "proper", hp_declared,
refining = NULL,
axis_refine = NULL) {
hp <- .joint_multi_hyperprior(
joint_grid, fn_sigma, fn_alpha, fn_phi, blocks = prepared,
families = hp_families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared)
res$log_hyperprior <- hp$lp
res$log_hyperprior_axes <- hp$axes
res$log_hyperprior_declined <- hp$declined
base <- !nzchar(.hyper_slice_home(refining, nrow(res$theta_grid)))
specs <- .joint_multi_measure_specs(
res$theta_grid[base, , drop = FALSE], hp$axes, joint_grid, axis_offsets,
B, fn_sigma, fn_alpha, fn_phi, copy_slab, copy_atom_mass,
logchol = hp_declared$logchol, axis_refine = axis_refine)
res$log_quad <- .hyper_log_quad_weights(res$theta_grid, specs,
refining = refining)
res$axis_support <- .hyper_grid_supports(res$theta_grid, specs,
refining = refining)
res <- .nl_record_unconverged(res)
res$weights <- .joint_integration_weights(res$log_marginal, dnode,
log_quad = res$log_quad)
list(res = res, specs = specs)
}
# Attach each hyperprior to the axis spec of the column it applies to, so the
# axis quadrature knows what the fold below will apply to that axis. That is
# what keeps the copy scale's declared point mass at its declared prior
# probability (`.hyper_atom_fold_scale()`).
.joint_multi_hp_specs <- function(specs, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi = NULL) {
if (is.null(specs)) return(specs)
if (is.null(fn_sigma) && is.null(fn_alpha) && is.null(fn_phi)) return(specs)
entries <- .joint_multi_hp_cols(joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi)
if (length(entries) == 0L) return(specs)
named <- colnames(joint_grid)[vapply(entries, function(e) e[["col"]],
integer(1))]
lapply(specs, function(sp) {
j <- match(sp$name, named)
if (is.na(j)) return(sp)
# The joint driver's hyperprior families are densities on the axis's
# natural scale.
sp$log_prior <- entries[[j]][["fn"]]
sp$log_prior_coord <- "natural"
sp
})
}
# Add the regularizing (sigma, alpha) hyperprior to a per-cell log-marginal
# vector for a multi-block `joint_grid`. Each applied axis is folded through
# the SAME `.joint_hp_vec_for_grids()` the single-block path uses, one column
# at a time, so the atom / continuum reading of a copy scale's zero is that
# axis's own. The per-axis contributions are summed before the one addition to
# `log_marginal`, so a grid carrying a single axis of each role reproduces the
# whole-view call to the bit. Single source of truth for the main dispatch and
# the Pareto-k re-evaluation.
.joint_multi_add_hp <- function(log_marginal, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi = NULL,
blocks = NULL, families = NULL,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
hyperprior = "proper", declared) {
log_marginal + .joint_multi_hyperprior(
joint_grid, fn_sigma, fn_alpha, fn_phi, blocks = blocks,
families = families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = declared)$lp
}
# The hyperprior record over a multi-block `joint_grid` (`.joint_hyperprior()`):
# the caller's densities on the roles they name, and with `blocks` supplied what
# `hyperprior` gives every other axis of each block. `declared` is what the
# fit's declared grid fixes (`.joint_multi_declared_axes()`).
.joint_multi_hyperprior <- function(joint_grid, fn_sigma, fn_alpha, fn_phi = NULL,
blocks = NULL, families = NULL,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
hyperprior = "proper", declared) {
.joint_hyperprior(joint_grid, blocks %||% list(), families,
user = list(sigma = fn_sigma, alpha = fn_alpha,
phi = fn_phi),
copy_atom_mass = copy_atom_mass,
hyperprior = if (is.null(blocks)) "flat" else hyperprior,
declared = declared)
}
# What a multi-block joint grid declares (`.hp_declare()`): the columns it
# integrates -- each block axis declared on more than one node, and each per-arm
# dispersion axis carrying a grid, named as the joint grid names them -- and each
# free-covariance block's design. Read off the declared per-block grids rather
# than any evaluated grid, so the tensor, a CCD design and every batch the driver
# evaluates share one answer.
.joint_multi_declared_axes <- function(block_grids, axis_names,
phi_axes = NULL) {
offsets <- cumsum(c(0L, vapply(block_grids, ncol, integer(1))))
latent <- character(0)
logchol <- list()
for (b in seq_along(block_grids)) {
if (!ncol(block_grids[[b]])) next
g <- block_grids[[b]]
colnames(g) <- as.character(axis_names)[(offsets[b] + 1L):offsets[b + 1L]]
d <- .hp_declare(g)
latent <- c(latent, d$axes)
logchol <- c(logchol, d$logchol)
}
phi <- if (length(phi_axes)) {
active <- vapply(phi_axes, function(v) length(unique(v)) > 1L, logical(1))
paste0("phi_", names(phi_axes)[active])
}
list(axes = c(latent, phi), logchol = logchol)
}
# Multi-block joint outer Pareto-k-hat. Builds the re-evaluation closure
# (round-trips a sampled user-facing `joint_grid` through the SAME kernel via
# the shared cpp-grid / phi / hyperprior helpers) and defers to the
# block-type-aware driver `.joint_pareto_k`. The whole importance batch is
# re-solved in one kernel call across `n_threads_outer`, so the independent
# re-solves run concurrently across cores rather than one-at-a-time using all
# inner threads. No tiling on the re-eval path (the IS batch is not a per-axis
# alpha lattice), so no tile-partition reconstruction is needed.
#
# The diagnostic solves are bounded: each is warm-started
# from the modal cell's converged latent mode (the bulk of the importance
# weight sits near it, so those draws converge in a few Newton steps) and
# capped at `.K_DIAG_MAX_ITER` iterations. A draw at an implausible
# hyperparameter where the inner Newton would otherwise stall to `max_iter`
# carries negligible importance weight, so the cap bounds its cost without
# moving the k-hat (converged draws keep their exact log-marginal; only the
# negligible-weight tail is truncated).
# Fit-scoped factory for the `.cpp_joint_multi` calls on this path (the CCD /
# adaptive mode-find, the node solves, the authoritative solve, the local-CCD
# refinement and the k-diagnostic re-solve). Everything invariant across those
# solves is captured once, so a per-fit setting cannot be honoured at one call
# site and silently dropped at another -- `control$hessian` was: it was parsed
# and match.arg-validated, then never passed on this path at all, so "psd" /
# "fisher" ran as LM/observed without a warning.
#
# The knobs a caller genuinely varies stay parameters. `max_iter` / `tol` /
# `inner_refresh` take fit-scoped defaults but are overridable, because the
# k-diagnostic deliberately re-solves with its own cheaper knobs.
.joint_multi_call_factory <- function(arms, cp, blocks_spec, axis_offsets, B,
n_threads, force_sparse, cell_coupling,
hessian_pd_mode, step_curvature_mode,
max_iter, tol, inner_refresh,
fixed_block_p = 0L,
fixed_block_constraints = NULL,
inner_sparse_override = 0L) {
force(arms); force(cp); force(blocks_spec); force(axis_offsets); force(B)
force(n_threads); force(force_sparse); force(cell_coupling)
force(inner_sparse_override)
force(hessian_pd_mode); force(step_curvature_mode)
force(max_iter); force(tol); force(inner_refresh)
force(fixed_block_p); force(fixed_block_constraints)
function(theta_mat,
x_init = NULL,
store_Q = FALSE,
phi_grid_per_arm = NULL,
n_threads_outer = 1L,
tile_ids = NULL,
tile_pilot_cells = NULL,
prune_tol = 0.0,
x_init_per_cell = NULL,
max_iter_ = NULL,
tol_ = NULL,
inner_refresh_ = NULL,
compute_skew = FALSE,
skew_idx = NULL,
# Per-cell fixed-effect covariance retention.
# Off by default so the mode-find, the Pareto-k re-solve batch and
# the skew probe -- none of which contribute an integrated cell --
# never pay for it; the grid solve and the local-CCD node solves
# switch it on.
fixed_block = FALSE,
# Subspace debias: the kernel-facing request
# list, off by default so no existing call site re-solves anything
# different.
debias = NULL,
# Corrected integrated Laplace: the
# kernel-facing request list, off by default for the same reason.
cila = NULL,
# Per-cell log hyperprior + log cell measure a screened solve
# ranks with, in `theta_mat` row order.
screen_log_offset = NULL,
# Return the screened surface without the full pass: a placement
# pilot's detecting grid only (`.nl_pilot_screen_only()`).
screen_only = FALSE) {
.cpp_joint_multi(
arms_list = arms,
copy_arms = as.integer(cp$copy_arms_zero),
copy_blocks = as.integer(cp$copy_blocks_zero),
blocks_spec = blocks_spec,
theta_grid = .joint_multi_cpp_grid(theta_mat, axis_offsets,
B, cp),
axis_offsets = axis_offsets,
max_iter = as.integer(max_iter_ %||% max_iter),
tol = as.numeric(tol_ %||% 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_ids,
tile_pilot_cells = 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_ %||% inner_refresh),
x_init_per_cell = x_init_per_cell,
compute_skew = isTRUE(compute_skew),
skew_idx = skew_idx,
fixed_block_p = if (isTRUE(fixed_block))
as.integer(fixed_block_p) else 0L,
fixed_block_constraints = fixed_block_constraints,
debias = debias,
cila = cila,
inner_sparse_override = as.integer(inner_sparse_override),
screen_log_offset = screen_log_offset,
screen_only = isTRUE(screen_only))
}
}
# The `kernel_fn(new_cells, warm_start, store_extras)` closure the refinement
# passes drive (`R/hyper_grid_refine.R`), over this driver's own kernel call: the
# multi-block counterpart of `.joint_make_kernel_fn()`. A refinement cell is a row
# of the assembled joint grid, so the per-arm dispersion columns it carries go
# through the same `phi_grid_per_arm` extraction the grid solve uses. The per-cell
# side data leaves only inside `extras`, so a call that asks for none requests
# neither the precision nor the fixed-effect block.
.joint_multi_make_kernel_fn <- function(call_kernel, arm_names, x_init_default,
store_Q = FALSE, screen_tol = 0,
screen_hp = NULL) {
function(new_cells, warm_start = NULL, store_extras = FALSE,
x_init_per_cell = NULL, screen = NULL) {
x_init <- if (!is.null(warm_start) && !is.null(warm_start$mode))
as.numeric(warm_start$mode) else x_init_default
scr <- .nl_refine_screen_offset(new_cells, screen, screen_tol, screen_hp)
res_x <- call_kernel(
new_cells, x_init = x_init,
store_Q = isTRUE(store_Q) && isTRUE(store_extras),
fixed_block = isTRUE(store_extras),
phi_grid_per_arm = .joint_multi_phi_per_arm(new_cells, arm_names),
x_init_per_cell = .joint_warm_start_per_cell(warm_start,
x_init_per_cell,
nrow(new_cells)),
prune_tol = if (is.null(scr)) 0 else screen_tol,
screen_log_offset = scr)
list(log_marginal = res_x$log_marginal,
extras = if (isTRUE(store_extras))
.joint_extras_from_res(res_x, nrow(new_cells)),
modes = res_x$modes)
}
}
# Multi-block counterpart of `.joint_attach_pareto_k_single()`
# (`R/nested_laplace_joint_pareto_k.R`). Same `diagnose_k = FALSE` behaviour:
# `pareto_k` / `pareto_k_is_ess` stay NA, and the placement mode the
# auto-recenter rescues read (`outer_mode_*`, `.joint_attach_placement()`) is
# attached either way.
.joint_attach_pareto_k_multi <- function(res, call_kernel,
axis_offsets, B, arm_names,
fn_sigma, fn_alpha, fn_phi = NULL,
max_iter, tol,
diagnose_k = TRUE,
diagnose_draws = .nl_diag("k_samples"),
n_threads_outer = 1L, proposal = NULL,
pareto_k_by_arm = FALSE,
k_bootstrap = .nl_diag("k_bootstrap"),
k_tail_points = NULL,
k_conf_bands = NULL,
placement_axes = character(0),
hp_blocks = NULL, hp_families = NULL,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
hyperprior = "proper", hp_declared) {
res$pareto_k <- NA_real_
res$pareto_k_is_ess <- NA_real_
res$pareto_k_scope <- "outer (hyperparameter) Gaussian proposal"
res$pareto_k_proposal_source <- NA_character_
res$pareto_k_declined <- NA_character_
# Regime read off stored weights: attached even with the diagnostic off.
res <- .joint_attach_pareto_k_regime(res)
warm_mode <- .joint_modal_mode(res)
modal_theta <- .joint_modal_theta(res)
k_max_iter <- min(as.integer(max_iter), .K_DIAG_MAX_ITER)
n_to <- as.integer(n_threads_outer)
knobs <- .kdiag_knobs()
# One kernel call over the importance batch. Shamanskii reuse makes the
# off-factor steps scatter grad-only (attacks the dominant per-iteration
# Hessian fill); loosened inner tol cuts intrinsic Newton steps; per-cell
# warm start (`x_init_per_cell`, an [S x n_x] matrix or NULL) starts each
# draw from its nearest grid mode in BOTH serial and parallel.
# The diagnostic re-solve runs with its own cheaper knobs (k_max_iter,
# knobs$tol / refresh), so its checkpoint fingerprint would differ from the
# main grid's; run it checkpoint-free so it neither collides with nor
# appends to the fit's checkpoint file.
# Importance draws are laid on each axis's unconstraining transform, the
# identity on a free-covariance block's columns, so they read the density in
# the column coordinates.
k_declared <- .hp_declare_on_columns(hp_declared)
solve_fn <- function(theta_mat, x_init_per_cell = NULL) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat,
x_init = warm_mode,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
n_threads_outer = n_to,
x_init_per_cell = x_init_per_cell,
max_iter_ = k_max_iter,
tol_ = max(knobs$tol, as.numeric(tol)),
inner_refresh_ = knobs$refresh))
.joint_multi_add_hp(r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = hp_blocks, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = k_declared)
}
# The placement mode-find at the fit's own inner settings (see
# `.joint_attach_pareto_k_single()`).
place_warm <- warm_mode
place_fn <- function(theta_mat) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat,
x_init = place_warm,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
n_threads_outer = n_to,
x_init_per_cell = .joint_rows_of(place_warm, nrow(theta_mat))))
lp <- .joint_multi_add_hp(r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = hp_blocks, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = k_declared)
if (is.matrix(r$modes)) attr(lp, "modes") <- r$modes
lp
}
place_set_warm <- function(mode) {
mode <- as.numeric(mode)
if (length(mode) && all(is.finite(mode))) place_warm <<- mode
}
if (!isTRUE(diagnose_k)) {
res <- .k_attach_declined(res, .k_decline("not_requested"))
return(.joint_attach_placement(res, place_fn, place_set_warm,
extra_axes = placement_axes))
}
# Per-cell warm start (nearest grid mode, serial + parallel) when modes are
# stored; else near-neighbour chain re-order; else plain broadcast
refit <- .joint_make_diag_refit(res, solve_fn, modal_theta, knobs)
arm_axes <- if (isTRUE(pareto_k_by_arm)) .joint_pareto_arm_axes(res) else NULL
kd <- .joint_pareto_k(res, refit, diagnose_draws, proposal = proposal,
arm_axes = arm_axes, k_bootstrap = k_bootstrap,
k_tail_points = k_tail_points, k_conf_bands = k_conf_bands)
res$pareto_k <- kd$pareto_k
res$pareto_k_is_ess <- kd$is_ess
res$pareto_k_proposal_source <- kd$proposal_source
res <- .k_attach_declined(res, kd)
res$pareto_k_mode_u <- kd$mode_u
res$pareto_k_cov_u <- kd$cov_u
res$pareto_k_axis_tags <- kd$axis_tags
res$pareto_k_axis_names <- kd$axis_names
res <- .joint_attach_pareto_k_regime(res, kd)
res <- .joint_attach_pareto_k_uncertainty(res, kd)
res <- .joint_attach_by_arm_k(res, kd)
.joint_attach_placement(res, place_fn, place_set_warm,
extra_axes = placement_axes)
}
# Inner-Laplace skewness diagnostic for the multi-block
# joint driver -- the multi-block analogue of
# .nlj_inner_skew_at_theta() (R/nested_laplace_joint.R). Reuses
# `call_kernel` (the SAME
# closure the outer Pareto-k diagnostic and the adaptive-grid refinement
# already call), re-dispatched at the fitted MAP grid cell with
# `compute_skew = TRUE`: one extra deterministic Newton solve, no importance
# sampling -- mirrors the single-block helper exactly except for the
# multi-block `call_kernel` signature (`x_init` instead of `warm_start`, and
# an explicit `phi_grid_per_arm` for the probed cell, since only the
# multi-block path carries per-arm dispersion axes).
#
# Default probe scope is every arm's fixed-effects coefficients, read from
# `res$arm_layout$beta_start` / `$p` (.joint_multi_layout(), already attached
# by the caller) -- the direct multi-block analogue of the single-block
# default ("every arm's beta").
#
# A coupled arm (`cell_coupling != "separable"` on that arm) has no per-obs
# oracle to score (build_joint_curvature3_fns excludes it) -- its indices
# come back NaN, not a silently wrong 0, exactly as for the single-block and
# single-arm paths.
.nlj_multi_inner_skew_at_theta <- function(res, call_kernel, arm_names,
skew_idx = NULL, compute = TRUE) {
res <- .inner_skew_decline(res, "not_requested")
if (!compute) return(res)
modal_theta <- .joint_modal_theta(res)
if (is.null(modal_theta)) return(.inner_skew_decline(res, "backend_unsupported"))
probe_idx <- skew_idx
if (is.null(probe_idx)) {
al <- res$arm_layout
if (is.null(al) || is.null(al$beta_start) || is.null(al$p)) {
return(.inner_skew_decline(res, "no_probe_indices"))
}
probe_idx <- unlist(lapply(seq_len(al$n_arms), function(k) {
if (al$p[k] <= 0L) return(integer(0))
(al$beta_start[k] + 1L):(al$beta_start[k] + al$p[k])
}))
}
probe_idx <- as.integer(probe_idx)
if (length(probe_idx) == 0L) return(.inner_skew_decline(res, "no_probe_indices"))
res <- .inner_skew_decline(res, "solve_failed")
warm_mode <- .joint_modal_mode(res)
theta_mat <- matrix(modal_theta, nrow = 1L,
dimnames = list(NULL, colnames(res$theta_grid)))
out <- tryCatch(
.joint_with_quiet_opts(call_kernel(
theta_mat, x_init = warm_mode,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
compute_skew = TRUE, skew_idx = probe_idx)),
error = function(e) NULL)
.inner_skew_attach_probe(res, out)
}
# Main multi-block joint dispatch. Mirrors the structure of the
# single-block path (build grid, call C++, post-process) but accepts a
# list-of-blocks `prior_list` and a `copy` spec that points at a
# specific block.
.joint_dispatch_multi <- function(responses, prior_list, copy,
phi_grid,
fn_sigma = NULL,
fn_alpha = NULL,
fn_phi = NULL,
max_iter, tol, n_threads,
x_init, verbose, store_Q,
keep_grid_hessians = TRUE,
n_threads_outer = 1L,
tile_warm = TRUE,
prune_tol = 0.0,
force_sparse = FALSE,
inner_sparse_override = 0L,
cell_coupling = "separable",
hessian_pd_mode = 0L,
step_curvature_mode = 0L,
diagnose_k = TRUE,
diagnose_draws = .nl_diag("k_samples"),
pareto_k_by_arm = FALSE,
pareto_k_threads = NULL,
k_bootstrap = .nl_diag("k_bootstrap"),
k_tail_points = NULL,
k_conf_bands = NULL,
diagnose_skew = TRUE,
skew_idx = NULL,
skew_correct = .nl_diag("skew_correct"),
subspace_debias = NULL,
cila = NULL,
inner_refresh = 1L,
integration = "auto",
within_cell = NULL,
local_ccd = NULL,
adaptive_cutoff = 10,
adaptive_stride = 2L,
adaptive_max_frac = 0.75,
adaptive_min_cells = 48,
copy_atom_mass = .TULPA_COPY_ATOM_MASS,
copy_slab = "exponential",
placement_axes = character(0),
hyperprior = "proper",
adaptive_grid = FALSE,
adaptive_grid_edge_thresh = 0.02,
adaptive_grid_max_passes = 1L,
var_of_means_consistency = TRUE,
axis_refine = NULL,
grid_auto = list(),
phi_auto = NULL,
timer = NULL) {
tm <- timer %||% .tulpa_timer()
integration <- match.arg(integration, c("auto", "ccd", "grid",
"grid_adaptive"))
within_cell <- .nl_within_cell_mode(within_cell)
n_arms <- length(responses)
arms <- lapply(seq_along(responses), function(k) {
a <- responses[[k]]
.normalise_joint_arm_multi(a, k)
})
arm_names <- names(responses) %||% paste0("arm", seq_along(responses))
cp <- .resolve_copy_multi(copy, responses, prior_list)
hp_families <- .joint_multi_hp_families(arms, arm_names)
# Per-block axis grids (with copy-block parameterisation if applicable).
B <- length(prior_list)
per_block <- lapply(seq_len(B), function(b) {
copy_pos <- if (cp$has_copy) match((b - 1L), cp$copy_blocks_zero)
else NA_integer_
is_copy <- !is.na(copy_pos)
alpha_grid_b <- if (is_copy) cp$alpha_grids[[copy_pos]] else numeric(0)
.joint_block_axis_grid(prior_list[[b]], is_copy, alpha_grid_b, b)
})
block_grids <- lapply(per_block, function(x) x$grid)
prepared <- lapply(per_block, function(x) x$prepared)
# Axis layout (independent of node placement: tensor or CCD).
axis_counts <- vapply(block_grids, ncol, integer(1))
axis_offsets <- as.integer(c(0L, cumsum(axis_counts)))
axis_names <- unlist(lapply(seq_along(block_grids), function(b) {
cn <- colnames(block_grids[[b]])
if (is.null(cn) || length(cn) == 0L) return(character(0))
paste0("b", b, ".", cn)
}))
d_axes <- as.integer(axis_offsets[B + 1L])
# Which axes the regularizing hyperpriors reach. Read off the axis layout,
# before the first inner solve, so a per-block entry naming a block that
# carries no such axis is refused rather than paid for.
hyperprior_axes <- .joint_check_multi_hyperpriors(
as.character(axis_names), axis_offsets, B, fn_sigma, fn_alpha)
if (isTRUE(verbose)) {
for (role in names(hyperprior_axes)) {
rec <- hyperprior_axes[[role]]
if (length(rec$axes) > 1L && identical(rec$scope, "every_block")) {
message("tulpa joint: prior_", role, " applies to ",
paste(rec$axes, collapse = ", "),
" (one spec, every block carrying the axis). ",
"Key it per block to regularize them differently.")
}
}
}
blocks_spec <- lapply(seq_along(prepared), function(b) {
.joint_block_spec_for_cpp(prepared[[b]], n_arms, b, arms = arms)
})
# Every `.cpp_joint_multi` call on this path goes through this closure, so
# the fit-scoped settings are supplied once instead of at each of the seven
# solve sites.
force_sparse <- .resolve_force_sparse(force_sparse, function() {
.joint_multi_layout(arms, prepared)$n_x
})
# Per-cell fixed-effect covariance retention. Both the
# block size and the field sum-to-zero groups are fixed by the latent layout,
# which is known before the first solve, so the request is built here and the
# kernel extracts each cell's block inside that cell's own solve
# instead of the grid keeping every precision.
fixed_layout <- .joint_fixed_layout(responses)
fb_layout <- .joint_multi_layout(arms, prepared)
fixed_block_p <- if (isTRUE(keep_grid_hessians))
as.integer(fixed_layout$n_fixed) else 0L
fixed_block_constraints <- .joint_constraint_cols(fb_layout, fb_layout$n_x)
call_kernel <- .joint_multi_call_factory(
arms = arms, cp = cp, blocks_spec = blocks_spec,
axis_offsets = axis_offsets, B = B, n_threads = n_threads,
force_sparse = force_sparse, cell_coupling = cell_coupling,
hessian_pd_mode = hessian_pd_mode,
step_curvature_mode = step_curvature_mode,
max_iter = max_iter, tol = tol, inner_refresh = inner_refresh,
fixed_block_p = fixed_block_p,
fixed_block_constraints = fixed_block_constraints,
inner_sparse_override = inner_sparse_override)
# Per-latent-axis fine grid values (sorted unique), shared by the CCD and the
# adaptive-lattice integrators to place / locate the outer nodes.
block_of_axis <- rep(seq_len(B), times = axis_counts)
col_within <- unlist(lapply(axis_counts, seq_len))
latent_axis_values <- lapply(seq_len(d_axes), function(j) {
sort(unique(as.numeric(block_grids[[block_of_axis[j]]][, col_within[j]])))
})
# Per-arm dispersion overrides on the outer grid. A phi axis varies
# independently of the latent-block axes, so it forms a Cartesian (tensor)
# product with the latent grid -- whether that latent grid is the dense
# tensor or a CCD design. The cross is applied once,
# below, to whichever latent grid was built.
phi_axes <- .normalise_phi_grid(phi_grid, arm_names)
has_phi <- !is.null(phi_axes) &&
any(vapply(phi_axes, length, integer(1)) > 0L)
hp_declared <- .joint_multi_declared_axes(block_grids, axis_names, phi_axes)
# How far each axis may be refined follows from where its nodes came from,
# the rule the single-block driver applies (`.joint_axis_refine_modes_by()`):
# nodes the caller stated are integrated more finely over exactly that range,
# nodes the engine placed may be followed out past their ends.
# `control$axis_refine` overrides either, per axis, and is resolved before
# the first kernel call so a mis-named axis is refused rather than reported
# after the grid has been solved.
refine_axes <- c(as.character(axis_names),
if (has_phi) paste0("phi_", names(phi_axes)[
vapply(phi_axes, length, integer(1)) > 0L]))
axis_refine_modes <- .joint_axis_refine_modes_by(
refine_axes,
function(a) .joint_multi_axis_is_stated(
a, prior_list, grid_auto, copy, as.integer(cp$copy_blocks_zero) + 1L,
phi_grid, arm_names, phi_auto),
user = .joint_check_axis_refine(axis_refine, refine_axes))
# Node placement for the LATENT axes. CCD integrates on a
# central composite design around the joint hyperparameter mode; it declines
# (-> tensor grid) for too few axes, an unguessable axis (CAR_proper rho_car
# / non-BYM2 rho), or a degenerate / ridged outer mode-find / Hessian
#. An active phi axis no longer disables CCD
#: the CCD rides the latent axes and the phi tensor
# crosses on top, with the CCD design weights replicated across phi cells.
# Axis threshold: "auto" (the default) engages CCD only at d >= 4, where the
# tensor blow-up bites hardest, and keeps the cheaper, ridge-robust tensor
# grid at d <= 3; explicit "ccd" lowers the threshold to d >= 3.
dnode <- NULL
tile_partition <- NULL
phi_grid_per_arm_list <- NULL
integration_used <- "grid"
joint_grid <- NULL
use_adaptive <- FALSE
adaptive_declined <- FALSE
adaptive_info <- NULL
# Mode-Hessian outer proposal. Set only when the CCD
# integrator engages: its Gaussian (centre `u_hat`, scale `L_scale` from the
# analytic curvature at the outer mode) drives the outer Pareto-k so the
# diagnostic stays defined when the hyperparameter posterior is sharp and the
# grid concentrates on ~1 cell (where the grid-weighted covariance collapses).
ccd_proposal <- NULL
ccd_requested <- .joint_ccd_engage(integration, d_axes)
# Why a CCD the caller asked for did not run. NA when
# nothing was declined, so `fit$integration == "grid"` is never the only
# trace of a request the engine turned down. An explicit
# `integration = "ccd"` below the axis threshold declines here; every other
# reason comes from .joint_ccd_grid()'s own vocabulary.
integration_declined <- NA_character_
if (identical(integration, "ccd") && !ccd_requested) {
integration_declined <- "axis_count"
}
# What placing the design cost, in inner solves. Carried on the placement's
# own return so a decline reports what it had already paid, and NULL where
# no placement ran -- a fit that placed nothing reports no cost rather than
# a zero, which would read as a placement that was free.
ccd_cost <- NULL
use_ccd <- ccd_requested
# Set only on the Cartesian-product branch below (CCD and the adaptive
# grid do not build a dense tensor, so the "reduce the grid" remedy the
# post-solve timing check reports does not apply to them).
n_cells <- NA_integer_
grid_warn_remedy <- NULL
grid_layout <- ""
if (use_ccd) {
axis_values <- latent_axis_values
# Warm latent mode for the mode-find / node solves: one solve at the
# box-centre coordinate, broadcast as x_init so each subsequent inner
# Newton converges in a few steps. Checkpoint/progress are stripped for
# every CCD probe so they neither pollute the fit's checkpoint file nor
# tick its progress bar.
ccd_warm <- x_init
centre_row <- matrix(vapply(axis_values, stats::median, numeric(1)),
nrow = 1L, dimnames = list(NULL, axis_names))
init_fit <- tryCatch(.joint_with_quiet_opts(call_kernel(
centre_row, x_init = x_init)),
error = function(e) NULL)
if (!is.null(init_fit) && is.matrix(init_fit$modes) &&
nrow(init_fit$modes) >= 1L &&
is.finite(init_fit$log_marginal[1L])) {
m1 <- as.numeric(init_fit$modes[1L, ])
if (all(is.finite(m1))) ccd_warm <- m1
}
# A CCD design is laid, and its mode found, in the grid's column
# coordinates, so the hyperprior it reads is the density in those.
ccd_declared <- .hp_declare_on_columns(hp_declared)
eval_logpost <- function(theta_mat) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat, x_init = ccd_warm,
n_threads_outer = n_threads_outer))
lp <- .joint_multi_add_hp(r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = prepared, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior,
declared = ccd_declared)
# Carry the inner latent modes so the CCD mode-find can advance the
# warm start per accepted point.
if (is.matrix(r$modes)) attr(lp, "modes") <- r$modes
lp
}
# Advance the inner warm start to the latent mode at each accepted
# mode-find point, so subsequent probes solve in a few Newton steps
# rather than cold from the box centre.
set_warm <- function(mode) {
mode <- as.numeric(mode)
if (length(mode) && all(is.finite(mode))) ccd_warm <<- mode
}
# The same atom mass the axis specs below declare, so the mixture the
# design splits into is weighted by the prior the grid is integrated
# against.
ccd <- .joint_ccd_grid(axis_names, axis_offsets, prepared, axis_values,
eval_logpost, verbose = verbose,
set_warm = set_warm,
atom_mass = copy_atom_mass)
ccd_cost <- ccd[c("modefind_evals", "modefind_budget",
"modefind_rounds", "modefind_seconds")]
if (is.null(ccd$grid)) {
use_ccd <- FALSE
integration_declined <- as.character(ccd$declined %||% "modefind_failed")
} else {
joint_grid <- ccd$grid
dnode <- ccd$dnode
integration_used <- "ccd"
hp_declared <- ccd_declared
# The CCD axes are exactly the leading latent-block columns of
# joint_grid (axis_names, in order); phi crosses as a separate
# tensor on top. Carry the mode-Hessian Gaussian over those axes for
# the outer Pareto-k.
#
# A design that split at a copy atom has no such Gaussian to carry.
# Its components are built in log alpha, while the diagnostic
# unconstrains the SAME grid on the identity coordinate -- the only
# one that represents the alpha = 0 cells the split places. One
# Gaussian does not describe a mixture either way, so the diagnostic
# keeps its grid-weighted proposal, which the split's spread of
# components supports.
if (isTRUE(ccd$atom_split)) {
ccd_proposal <- NULL
} else {
ccd_proposal <- list(u_hat = ccd$u_hat,
L_scale = ccd$L_scale,
tags = ccd$tags,
cols = seq_len(d_axes))
}
}
}
# Adaptive lattice (nested_laplace_joint_adaptive.R): the low-dimensional
# companion to the CCD. It evaluates a coarse subsample of the full outer
# lattice (latent block axes AND phi axes together) to locate the posterior
# mass, then floods outward on the fine lattice keeping only the cells within
# a log-density cutoff of the peak -- so the far tail, which every dense cell
# still pays a full inner solve for, is never placed. The selected cells are
# a strict subset of the dense tensor at uniform lattice spacing, so the main
# kernel call below evaluates them with the full contract (store_Q, phi, tile
# warm start) and the integration weight stays the plain softmax (dnode NULL);
# the flood solves here are cheap grid selection (store_Q off, quiet). Declines
# to the dense tensor on a degenerate lattice or a diffuse posterior whose kept
# region rivals the dense grid.
if (!use_ccd && .joint_adaptive_engage(integration, d_axes)) {
if (has_phi) {
active_ad <- phi_axes[vapply(phi_axes, length, integer(1)) > 0L]
phi_axis_vals <- lapply(active_ad, function(v)
sort(unique(as.numeric(v))))
phi_cols_ad <- paste0("phi_", names(active_ad))
} else {
phi_axis_vals <- list()
phi_cols_ad <- character(0)
}
ad_axis_values <- c(latent_axis_values, phi_axis_vals)
ad_col_names <- c(axis_names, phi_cols_ad)
# One warm centre solve, broadcast as x_init so each flood-selection solve
# converges in a few Newton steps rather than cold from the box centre.
ad_warm <- x_init
centre_theta <- matrix(
vapply(ad_axis_values, stats::median, numeric(1)), nrow = 1L,
dimnames = list(NULL, ad_col_names))
init_ad <- tryCatch(.joint_with_quiet_opts(call_kernel(
centre_theta, x_init = x_init,
phi_grid_per_arm = .joint_multi_phi_per_arm(centre_theta, arm_names))),
error = function(e) NULL)
if (!is.null(init_ad) && is.matrix(init_ad$modes) &&
nrow(init_ad$modes) >= 1L && is.finite(init_ad$log_marginal[1L])) {
m1 <- as.numeric(init_ad$modes[1L, ])
if (all(is.finite(m1))) ad_warm <- m1
}
# Grid-selection evaluator: physical theta (latent + phi columns) -> the
# inner joint log-marginal with the hyperprior baked in. store_Q off and
# quiet, since these solves only rank cells; the authoritative solve (with
# store_Q) is the main kernel call over the selected grid.
eval_theta_ad <- function(theta_mat) {
r <- .joint_with_quiet_opts(call_kernel(
theta_mat, x_init = ad_warm,
phi_grid_per_arm = .joint_multi_phi_per_arm(theta_mat, arm_names),
n_threads_outer = n_threads_outer))
list(log_marginal = .joint_multi_add_hp(
r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = prepared, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared),
modes = NULL)
}
ad <- .joint_adaptive_grid(ad_axis_values, ad_col_names, eval_theta_ad,
cutoff = adaptive_cutoff,
stride = adaptive_stride,
max_frac = adaptive_max_frac,
min_dense = adaptive_min_cells,
verbose = verbose)
adaptive_declined <- !is.null(ad$declined)
if (adaptive_declined) {
integration_declined <- ad$declined
} else {
joint_grid <- ad$grid
integration_used <- "grid_adaptive"
use_adaptive <- TRUE
adaptive_info <- ad$info
# phi columns are already part of the adaptive grid; extract them for
# the C++ call and skip the phi tensor cross below.
phi_grid_per_arm_list <- .joint_multi_phi_per_arm(joint_grid, arm_names)
}
}
if (!use_ccd && !use_adaptive) {
# Cartesian product of per-block axis grids.
# The product over zero blocks is one cell with no latent axis, so a
# model whose only hyperparameters are dispersions is solved once per
# dispersion cell crossed on top.
row_counts <- vapply(block_grids, nrow, integer(1))
idx <- if (B > 0L) do.call(expand.grid, lapply(row_counts, seq_len))
else data.frame(row.names = 1L)
# Every latent cell is solved once per dispersion cell crossed on top
# below, so the count the timing warning reads is the crossed one: that
# is the number of inner solves the fit pays.
phi_active <- if (has_phi) phi_axes[lengths(phi_axes) > 0L]
n_cells <- nrow(idx) * prod(lengths(phi_active))
# The CCD and the adaptive lattice are remedies for the cell count only
# while they are still available. Once this fit has already declined
# one, advising the caller to set the option they set is noise, so the
# advice names the decline instead. The adaptive lattice is named first:
# it integrates the same lattice, where the CCD is a different design.
alt <- c(
if (!adaptive_declined)
"\"grid_adaptive\" (the same lattice, solved where the posterior mass sits)",
if (adaptive_declined || is.na(integration_declined)) "\"ccd\"")
declined <- if (!is.na(integration_declined))
sprintf(if (adaptive_declined)
"the adaptive lattice declined to the dense tensor (%s)"
else "CCD integration declined (%s)",
integration_declined)
grid_warn_remedy <- paste0(
if (length(phi_active)) "Reduce per-block or phi grid sizes"
else "Reduce per-block grid sizes",
if (length(alt))
paste0(" or set control$integration = ",
paste(alt, collapse = " or ")),
if (length(declined))
paste0("; ", paste(declined, collapse = "; ")),
".")
grid_layout <- .nl_grid_crossing(block_grids, phi_active)
.nl_check_grid_cap(n_cells, .nl_max_grid_cells(), grid_warn_remedy,
grid_layout)
.nl_dense_grid_warn(n_cells, grid_warn_remedy, grid_layout)
joint_grid <- if (B > 0L) do.call(cbind, lapply(seq_along(block_grids),
function(b) block_grids[[b]][idx[[b]], , drop = FALSE]))
else matrix(numeric(0), nrow = 1L, ncol = 0L)
if (ncol(joint_grid) > 0L) colnames(joint_grid) <- axis_names
}
# Latent node / cell count, before the phi tensor crosses on top -- the CCD
# node count or the tensor latent-cell count, for the verbose announcement.
n_latent_cells <- nrow(joint_grid)
# phi_grid (per-arm dispersion overrides on the outer grid): each phi axis
# varies independently of the latent-block grid, so it forms a Cartesian
# product with the latent grid -- the dense tensor or the CCD design alike
#. On the CCD path the per-node design weights `dnode`
# are replicated across the phi cells so each (latent node, phi cell) keeps
# the latent node's quadrature weight; phi rides as a uniform tensor axis,
# exactly as on the dense path. A phi axis tied to a latent-block axis is
# not expressible here; fold it into the relevant block's own axis grid.
if (has_phi && !use_adaptive) {
crossed <- .joint_multi_cross_phi(joint_grid, phi_axes, axis_names)
joint_grid <- crossed$grid
# phi columns don't belong to any latent block; the C++ entry consumes
# them via phi_grid_per_arm, not the block axes.
phi_grid_per_arm_list <- .joint_multi_phi_per_arm(joint_grid, arm_names)
# Replicate the CCD design weights across the crossed phi cells (NULL on
# the dense path, where weights are the uniform-cell softmax).
if (!is.null(dnode)) dnode <- dnode[crossed$latent_row]
}
# Announce the engaged outer integrator at selection time:
# the single authoritative line, tied to integration_used, before the heavy
# inner solves -- so a CCD engaged silently by axis count (or declined back
# to the tensor on a ridge) is visible without reading fit$...$integration.
if (isTRUE(verbose) && !use_adaptive) {
n_phi_cells <- as.integer(round(nrow(joint_grid) / n_latent_cells))
.joint_announce_integration(
integration_used, d_axes, n_latent_cells, n_phi_cells,
nrow(joint_grid),
declined = if (adaptive_declined) "adaptive lattice"
else if (ccd_requested && !identical(integration_used, "ccd"))
"CCD")
}
# Tile partition for the three-tier warm-start.
# Tensor-grid-only: the scattered CCD design is not a
# per-axis lattice the alpha-tile reuse relies on. Tile axis = every
# joint_grid column EXCEPT the copy block's alpha column. Built from the
# *user-facing* (sigma, alpha) grid (before sigma_pos materialisation) so
# the partition reflects what is constant across alpha at fixed (sigma, rho,
# ... other-block axes).
if (!use_ccd && isTRUE(tile_warm) && cp$has_copy &&
length(cp$copy_blocks_zero) == 1L &&
as.integer(n_threads_outer) > 1L) {
b_copy_R <- cp$copy_blocks_zero[[1L]] + 1L
cols_bc <- (axis_offsets[b_copy_R] + 1L):axis_offsets[b_copy_R + 1L]
bare_bc <- sub("^b[0-9]+\\.", "", colnames(joint_grid)[cols_bc])
i_alpha <- match("alpha", bare_bc)
if (!is.na(i_alpha)) {
alpha_col_idx <- cols_bc[i_alpha]
non_alpha_idx <- setdiff(seq_len(ncol(joint_grid)), alpha_col_idx)
tile_partition <- .joint_compute_tile_partition(
joint_grid[, non_alpha_idx, drop = FALSE],
as.numeric(joint_grid[, alpha_col_idx]),
nrow(joint_grid)
)
}
}
# The regularizing hyperprior, folded into log_marginal at the kernel-call
# boundary below, and the base tensor's cell measure: known from the grid
# alone, so a screened solve ranks its cells with both.
hp_base <- .joint_multi_hyperprior(
joint_grid, fn_sigma, fn_alpha, fn_phi, blocks = prepared,
families = hp_families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared)
screen_offset <- if (as.numeric(prune_tol) > 0)
.nl_screen_log_offset(joint_grid, list(hp_base),
specs = .joint_multi_measure_specs(
joint_grid, hp_base$axes, joint_grid,
axis_offsets, B, fn_sigma, fn_alpha, fn_phi,
copy_slab, copy_atom_mass,
logchol = hp_declared$logchol))
# A placement pilot's detecting grid is read off the screen alone
# (`.nl_pilot_control()`); the full-grid fallback below never is.
pilot_screen <- isTRUE(getOption("tulpa.nl_pilot_screen", FALSE))
call_kernel_with_tol <- function(tol_prune) {
call_kernel(
joint_grid,
x_init = x_init,
store_Q = store_Q,
fixed_block = TRUE,
phi_grid_per_arm = phi_grid_per_arm_list,
n_threads_outer = n_threads_outer,
tile_ids = tile_partition$tile_ids,
tile_pilot_cells = tile_partition$tile_pilot_cells,
prune_tol = tol_prune,
screen_log_offset = if (tol_prune > 0) screen_offset,
screen_only = pilot_screen && tol_prune > 0)
}
tm$mark("setup")
grid_solve_start <- proc.time()[["elapsed"]]
res <- .joint_main_grid_solve(function() call_kernel_with_tol(prune_tol))
# Safety gate: fall back to the full grid when the
# cheap-pass ranking is unreliable rather than silently returning a
# pruned answer.
if (as.numeric(prune_tol) > 0) {
res <- .joint_prune_safety_gate(
res, resolve_full = function() call_kernel_with_tol(0.0))
}
if (!is.null(grid_warn_remedy)) {
.nl_multi_grid_warn(proc.time()[["elapsed"]] - grid_solve_start,
n_cells, grid_warn_remedy, grid_layout)
}
tm$mark("grid")
# Bake the regularizing hyperprior on (sigma, alpha) into log_marginal
#. Multi-block has no in-package refinement passes,
# so one apply at the kernel-call boundary suffices.
res <- .nl_fold_hyperprior(res, list(hp_base))
# The base tensor's absolute cell measure, read before local CCD refinement
# replaces cells by design-weighted clouds whose shares are relative to it.
base_grid <- joint_grid
# Local CCD refinement: replace a few high-weight, mutually
# non-adjacent tensor cells with small curvature-aware node clouds so a coarse
# base grid resolves the sharply-peaked directions without the k^d tensor
# blow-up. Engages only on the tensor path (the finite-difference curvature
# stencil needs axis neighbours), at >= 4 transformable latent axes, with no
# active phi grid. The node solves reuse the kernel through a quiet evaluator
# (checkpoint / progress stripped) warm-started per cell; refined cells carry
# partition-of-unity design weights so the total integration weight is
# conserved (no double-count).
local_ccd_info <- NULL
weight_kind <- NULL
# Local-CCD refinement changes the outer grid but not the stored per-cell Q
# (nor the tile partition / phi tensor), so combining it with store_Q would
# leave a Q that no longer aligns with the refined grid and crash
# tulpa_posterior_draws. Keeping Q valid for draws takes precedence: skip
# the refinement when store_Q is on.
if (!is.null(local_ccd) && isTRUE(store_Q) && !use_ccd && !use_adaptive &&
!has_phi && .joint_local_ccd_engage(d_axes) && isTRUE(verbose)) {
message("tulpa joint: local-CCD refinement skipped (store_Q = TRUE); ",
"keeping the per-cell Q aligned with the grid for posterior draws.")
}
if (!is.null(local_ccd) && !isTRUE(store_Q) && !use_ccd && !use_adaptive &&
!has_phi && .joint_local_ccd_engage(d_axes)) {
lc <- if (is.list(local_ccd)) local_ccd else list()
lc_max_cells <- as.integer(lc$max_cells %||% 8L)
lc_f0 <- as.numeric(lc$f0 %||% 1.1)
lc_skew_max <- as.numeric(lc$skew_max %||% .nl_diag("gamma3_ok"))
lc_rank <- match.arg(lc$rank %||% "weight", .LCCD_RANKS)
tags_lc <- .joint_ccd_axis_tags(axis_names, axis_offsets, prepared)
if (!is.null(tags_lc)) {
eval_nodes <- function(theta_mat, warm) {
xi <- if (!is.null(warm) && length(warm) && all(is.finite(warm)))
as.numeric(warm) else x_init
r <- .joint_with_quiet_opts(call_kernel(
theta_mat, x_init = xi,
fixed_block = (fixed_block_p > 0L),
n_threads_outer = n_threads_outer))
list(log_marginal = .joint_multi_add_hp(
r$log_marginal, theta_mat, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi,
blocks = prepared, families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared),
modes = if (is.matrix(r$modes)) r$modes else NULL,
cov_blocks = r$cov_block_per_grid)
}
ref <- .joint_local_ccd_refine(
joint_grid = joint_grid, log_marginal = res$log_marginal,
modes = res$modes, cov_blocks = res$cov_block_per_grid,
dnode = dnode, latent_axes = axis_names,
tags = tags_lc, eval_nodes = eval_nodes,
max_cells = lc_max_cells, f0 = lc_f0, verbose = verbose,
skew_max = lc_skew_max, rank = lc_rank)
if (!is.null(ref)) {
joint_grid <- ref$joint_grid
res$log_marginal <- ref$log_marginal
res$modes <- ref$modes
dnode <- ref$dnode
local_ccd_info <- ref$info
weight_kind <- ref$weight_kind
# The node solves carry their own fixed-effect block, so the
# refined grid keeps the retention aligned with the weights it
# is integrated against. The per-cell
# precision is not carried, which is why refinement engages only
# when store_Q is off -- the gate above.
res$cov_block_per_grid <- ref$cov_blocks
# The per-cell linear predictor is not carried either, and
# unlike the precision it is not gated: it still indexes the
# grid the fit had BEFORE this pass. Dropped rather than read
# against the new cells, so the grid-mixture predictive read
# falls back to the fit's coefficients instead of drawing from
# cells that are no longer the fit's (gcol33/tulpa#850).
res$fitted_eta <- NULL
res$fitted_eta_var <- NULL
}
}
tm$mark("grid")
}
# Outer-grid refinement on the tensor grid, the passes the single-block
# driver runs (`.joint_refine_outer_grid()`): an axis whose marginal has
# collapsed onto too few nodes to carry a spread is bisected in the modal
# cell's row until it resolves, and the boundary / interior pass follows the
# posterior past an engine-placed axis's ends when asked for. A CCD design,
# the adaptive lattice and a locally refined grid are not lattices a node can
# be inserted into, so they keep the nodes they were laid on.
joint_grid_init <- joint_grid
refining_axis <- NULL
refine_on <- !use_ccd && !use_adaptive && is.null(local_ccd_info) &&
(isTRUE(adaptive_grid) || isTRUE(var_of_means_consistency))
refine_hp_record <- function(cells)
.joint_multi_hyperprior(
cells, fn_sigma, fn_alpha, fn_phi, blocks = prepared,
families = hp_families, copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior, declared = hp_declared)
refine_hp <- function(cells) refine_hp_record(cells)$lp
refine_kernel <- .joint_multi_make_kernel_fn(call_kernel, arm_names, x_init,
store_Q,
screen_tol = as.numeric(prune_tol),
screen_hp = refine_hp_record)
grid_lm <- res$log_marginal
grid_extras <- .joint_init_extras_from_res(res)
n_added <- 0L
if (refine_on) {
refine_specs <- .joint_multi_measure_specs(
joint_grid, hp_base$axes, joint_grid, axis_offsets, B,
fn_sigma, fn_alpha, fn_phi, copy_slab, copy_atom_mass,
logchol = hp_declared$logchol, axis_refine = axis_refine_modes)
ref <- .joint_refine_outer_grid(
joint_grid, grid_lm, grid_extras, refine_specs, refine_kernel,
refine_hp,
adaptive_grid = adaptive_grid,
edge_thresh = adaptive_grid_edge_thresh,
max_passes = adaptive_grid_max_passes,
consistency = var_of_means_consistency,
axis_modes = .nl_outer_mode_axes(
getOption("tulpa.nl_outer_mode")))
refining_axis <- ref$refining_axis
joint_grid <- ref$theta_grid
grid_lm <- ref$log_marginal
grid_extras <- ref$extras
n_added <- ref$n_added
adaptive_info <- ref$adaptive_info
if (isTRUE(var_of_means_consistency)) {
res$var_of_means_consistency_info <- ref$consistency_info
}
}
# A cell whose inner Newton stopped short of a mode, on any node set this
# driver lays, is solved again from its nearest converged neighbour's mode
# before the grid is read.
rs <- .joint_resolve_unconverged(joint_grid, grid_lm, grid_extras,
refine_kernel, refine_hp)
if (n_added > 0L || rs$info$n_resolved > 0L) {
res <- .joint_glue_extras_to_res(res, joint_grid, rs$log_marginal,
rs$extras, refining_axis)
}
res$refining_axis <- refining_axis
res$nonconverged_resolve <- rs$info
tm$mark("grid")
res$theta_grid <- joint_grid
res$theta_names <- colnames(joint_grid)
res$axis_offsets <- axis_offsets
# Which grid axes each regularizing hyperprior was folded onto, and whether
# the caller stated it once (every block carrying that axis) or per block.
# NULL where the fit carries none, so a flat fit reports nothing rather
# than an empty record.
res$hyperprior_axes <- hyperprior_axes
res$integration <- integration_used
# What the caller asked for, and why it did not run. `integration` alone
# names the integrator that RAN, so a fit reporting "grid" cannot say
# whether the tensor grid was chosen or fallen back to.
# `.nl_node_support()` keys the interval construction off `integration`, so
# the difference is one a consumer can act on.
res$integration_requested <- integration
res$integration_declined <- integration_declined
res$ccd_modefind_evals <- ccd_cost$modefind_evals
res$ccd_modefind_budget <- ccd_cost$modefind_budget
res$ccd_modefind_rounds <- ccd_cost$modefind_rounds
res$ccd_modefind_seconds <- ccd_cost$modefind_seconds
# Integration weights fold in the CCD design weights (`dnode`) and, on the
# tensor grid, the per-cell prior mass of the node each cell sits at. A CCD
# design carries its own volume in `dnode`, so `log_quad` applies to the
# tensor path alone.
integ <- .joint_multi_attach_integration(
res, joint_grid, axis_offsets, B, prepared, hp_families,
fn_sigma, fn_alpha, fn_phi, copy_slab, copy_atom_mass, dnode = dnode,
hyperprior = hyperprior, hp_declared = hp_declared,
refining = refining_axis,
axis_refine = if (refine_on) axis_refine_modes)
res <- .nl_prune_merged_record(integ$res)
multi_specs <- integ$specs
if (refine_on) {
res$axis_span <- .joint_axis_span(joint_grid_init, res$theta_grid,
multi_specs, refining = refining_axis)
}
# The outer design weight each cell carries, kept beside the integration
# weight it was folded into: absent on a tensor base (uniform cell weight),
# the CCD design weights on a global CCD, the partition-of-unity shares of
# each replacement cloud on a locally refined grid. Recovering it from
# `weights` afterwards is a division by exp(log_marginal), which loses its
# scale and is undefined on a cell whose inner solve returned no finite
# marginal, so the grid state a weight-construction experiment reads is
# recorded rather than reconstructed.
res$dnode <- dnode
is_ccd <- identical(integration_used, "ccd")
# The evidence is read under the same measure the weights take: the cell
# masses on a tensor grid, the partition-of-unity design shares on a
# locally refined one. A global CCD's design weights reproduce moments and
# carry no volume, so no evidence is read off them.
if (is_ccd) {
res$log_evidence <- NA_real_
res$log_evidence_declined <- "moment_rule_design"
} else if (is.null(dnode)) {
res <- .nl_attach_evidence(res, res$theta_grid, multi_specs)
} else {
# A refined cloud's design shares split the base cell it replaced, so
# each node's absolute measure is its share of that cell's. The base
# tensor is evenly spaced wherever refinement engages, so one cell's
# measure serves every node; an uneven base has no single value to share.
base_specs <- .joint_axis_specs_from_grid(
base_grid, copy_slab = copy_slab,
folded_axes = res$log_hyperprior_axes)
base_lq <- .hyper_log_quad_weights(base_grid, base_specs, absolute = TRUE)
even <- all(is.finite(base_lq)) && diff(range(base_lq)) < 1e-8
res <- .nl_attach_evidence(
res, res$theta_grid, multi_specs,
log_measure = if (even) log(pmax(dnode, 0)) + base_lq[1L]
else rep(NA_real_, length(dnode)))
}
# What kind of weight each cell carries. `integration` names the integrator
# that ran, which describes a homogeneous support: a tensor cell holds the
# mass of its own cell, a CCD node holds a design weight. A locally refined
# grid is the one case with both, so it says so per cell rather than leaving
# a consumer to read one kind off `integration`.
res$weight_kind <- weight_kind %||%
rep(if (is_ccd) "design" else "mass", nrow(joint_grid))
if (!is.null(local_ccd_info)) {
# The share of the posterior that sits on design-weighted nodes: the
# part of the support whose cumulative weight is not a CDF, and so the
# bound on how much mass the per-axis quantile can misplace. It is also
# confined to the refined cells' own Voronoi boxes, so the displacement
# is at most one base-grid cell wide.
local_ccd_info$design_mass <-
sum(res$weights[res$weight_kind == "design"], na.rm = TRUE)
}
# Local CCD refinement also leaves a design-weighted (scattered) grid, so it
# takes the same moment / SD path as the global CCD: weighted moments over the
# design, no per-axis lattice SD refit. The adaptive lattice is a mass-
# concentrated SUBSET of the tensor (gaps in the far tail), so its per-axis
# levels are not the regular full lattice the 3-point profile SD refit needs;
# and it is dense near the peak, so the weighted var-of-means over the kept
# cells is already the calibrated SD. It therefore takes the same weighted-
# moment path, not the lattice refit.
# A refinement whose every candidate cell was put back as a mass atom
# leaves a plain tensor grid, so it takes the density path
# its cells carry rather than the design-weighted one it asked for.
is_design_weighted <- is_ccd || use_adaptive ||
(!is.null(local_ccd_info) && local_ccd_info$n_design_nodes > 0L)
# The median and interval per axis are read off whatever node set the
# integrator left. The global CCD is a MOMENT RULE -- its nodes reproduce the
# integrand's moments and carry no mass -- so a cumulative sum over them is
# not a CDF and the interval comes from the moments instead. The adaptive
# lattice keeps a cell per unit of mass, so its cumulative weight is a CDF.
# A locally CCD-refined grid is neither: it is part cell masses and part
# design, and says so rather than passing for a density grid. What it
# REPORTS is still the weighted quantile -- that measured best against a
# converged reference in both design_mass regimes, see
# `.nl_summary_quantile` -- so naming the support changes what the fit can
# tell a caller, not the numbers.
prov <- .nl_interval_provenance(integration_used, res$weight_kind,
res$weights)
res$theta_interval_read <- prov$read
res$theta_interval_design_mass <- prov$design_mass
res <- .joint_posterior_moments_multi(
res, prepared, axis_offsets, joint_grid, cp,
int_weights = if (is_design_weighted) res$weights else NULL,
support = prov$read, within = within_cell)
res$arm_layout <- .joint_multi_layout(arms, prepared)
res$blocks <- prepared
res$prior <- prior_list
res$responses <- responses
res <- .joint_attach_observations(res, responses)
res$copy <- copy
res$cell_coupling <- cell_coupling
res$local_ccd_info <- local_ccd_info
res$adaptive_grid_info <- adaptive_info
# Every axis's declared nodes, whatever subset of the lattice the adaptive
# flood, a pilot thinning or the prune screen kept: the placement mode-find
# searches an axis over these (`.nl_placement_mode()`).
phi_declared <- if (has_phi)
phi_axes[vapply(phi_axes, length, integer(1)) > 0L] else list()
res$outer_axis_nodes <- stats::setNames(
c(latent_axis_values,
lapply(phi_declared, function(v) sort(unique(as.numeric(v))))),
c(axis_names, if (length(phi_declared))
paste0("phi_", names(phi_declared))))
tm$mark("postproc")
# Outer Pareto-k-hat: re-evaluate the inner joint marginal at sampled
# hyperparameters via the shared cpp-grid / phi / hyperprior helpers and
# PSIS-smooth. Declines (NA -> quad-ESS) when a block carries an axis with
# unguessable support (CAR_proper's rho_car, a non-BYM2 rho, ...).
# The diagnostic's importance batch runs across `pareto_k_threads` (resolved
# by the top-level caller from control$k_threads; NULL only on a direct call,
# where it follows this fit's own thread grant)..
k_to <- pareto_k_threads %||%
.tulpa_pareto_k_threads(n_threads_outer, n_threads, diagnose_draws, NULL)
res <- .joint_attach_pareto_k_multi(res, call_kernel,
axis_offsets, B, arm_names,
fn_sigma, fn_alpha, fn_phi,
max_iter, tol,
diagnose_k = diagnose_k,
diagnose_draws = diagnose_draws,
n_threads_outer = k_to,
proposal = ccd_proposal,
pareto_k_by_arm = pareto_k_by_arm,
k_bootstrap = k_bootstrap,
k_tail_points = k_tail_points,
k_conf_bands = k_conf_bands,
placement_axes = placement_axes,
hp_blocks = prepared,
hp_families = hp_families,
copy_atom_mass = copy_atom_mass,
hyperprior = hyperprior,
hp_declared = hp_declared)
res <- .nlj_multi_inner_skew_at_theta(res, call_kernel, arm_names,
skew_idx, compute = diagnose_skew)
fixed <- .joint_fixed_layout(responses)
res <- .nl_skew_correction_attach(res, fixed$n_fixed, skew_correct)
tm$mark("diagnostics")
# Per-cell fixed-effect mode + precision for the grid marginalization.
# Local-CCD refinement invalidates the stored cells, so a fit it engaged on
# records that reason instead.
res <- .joint_finalize_grid_fixed(res, fixed$n_fixed, keep_grid_hessians)
# Subspace debias, over the settled grid, through the
# same `call_kernel` the skew probe and the Pareto-k batch re-dispatch with.
res <- .nl_subspace_debias_attach(
res, subspace_debias,
redispatch = function(req) .joint_with_quiet_opts(call_kernel(
res$theta_grid,
phi_grid_per_arm = .joint_multi_phi_per_arm(res$theta_grid,
arm_names),
debias = req)),
p_fixed = fixed$n_fixed, beta_names = fixed$names)
# Corrected integrated Laplace, over the same settled
# grid and through the same `call_kernel`.
res <- .nl_cila_attach(
res, cila,
redispatch = function(req) .joint_with_quiet_opts(call_kernel(
res$theta_grid,
phi_grid_per_arm = .joint_multi_phi_per_arm(res$theta_grid,
arm_names),
cila = req)),
p_fixed = fixed$n_fixed, beta_names = fixed$names,
remoments = function(r) {
r <- .joint_posterior_moments_multi(
r, prepared, axis_offsets, joint_grid, cp,
int_weights = if (is_design_weighted) r$weights else NULL,
support = prov$read, within = within_cell)
r
})
res$timing <- tm$timing()
res <- .joint_attach_diagnose_cost(res, diagnose_k, diagnose_draws)
.finalize_fit(res, backend = "nested_laplace_joint",
n_fixed = fixed$n_fixed, fixed_names = fixed$names,
extra_class = c("tulpa_nested_laplace_joint_multi",
"tulpa_nested_laplace_joint",
"tulpa_nested_laplace", "list"))
}
# Fixed-effect layout of a joint fit: the latent stacks all arm betas as a
# contiguous prefix (see .joint_multi_layout), so the generic accessors can
# slice coef/summary to the fixed effects instead of the full latent vector.
.joint_fixed_layout <- function(responses) {
arm_names <- names(responses) %||% paste0("arm", seq_along(responses))
nm <- unlist(lapply(seq_along(responses), function(k) {
X <- responses[[k]]$X
cn <- colnames(X) %||% paste0("beta", seq_len(ncol(X)))
cn[!nzchar(cn)] <- paste0("beta", which(!nzchar(cn)))
paste(arm_names[k], cn, sep = ".")
}), use.names = FALSE)
list(n_fixed = length(nm), names = nm)
}
# Latent-vector layout for the multi-block joint result. Mirrors the
# single-block `.joint_layout()` but generalised over `prepared` blocks:
# per-arm beta, per-arm RE, then each prepared block in order. For
# back-compat with single-block consumers (e.g. tulpaObs cover-hurdle's
# `.joint_field_at_obs_copy_multi`), the first spatial-like block also
# emits `phi_start` / `theta_start` aliases:
#
# * BYM2 -> phi_start, theta_start (length-2 sub-block).
# * ICAR / CAR_proper -> phi_start (length-1).
#
# Non-spatial blocks expose `block_start[b]` so callers can index into
# any block by ordinal position.
.joint_multi_layout <- function(arms, prepared) {
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]
}
B <- length(prepared)
block_start <- integer(B)
block_size <- integer(B)
phi_start <- NULL
theta_start <- NULL
# Per-spatial-block latent offset (0-based) of the structured field, in
# block order: ICAR / CAR_proper at the block start, BYM2 at its phi
# sub-block start. `phi_start` (back-compat) is field_starts[1].
field_starts <- integer(0)
field_block_types <- character(0)
field_n_fields <- integer(0)
for (b in seq_len(B)) {
block_start[b] <- cur
type <- tolower(prepared[[b]]$type %||% "")
# Latent size per block. BYM2 is length-2 (phi ICAR + theta IID)
# on n_spatial_units; all other blocks are length-1 on the
# block's `n_spatial_units` / `n_times` / `n_units`.
n_units <- prepared[[b]]$n_spatial_units %||%
prepared[[b]]$n_times %||%
prepared[[b]]$n_units %||%
prepared[[b]]$n_groups %||%
prepared[[b]]$n_latent %||%
prepared[[b]]$n_mesh %||%
prepared[[b]]$m_total %||% 0L
n_units <- as.integer(n_units)
if (type == "bym2") {
block_size[b] <- 2L * n_units
field_starts <- c(field_starts, cur)
field_block_types <- c(field_block_types, "bym2")
field_n_fields <- c(field_n_fields, 1L)
if (is.null(phi_start)) {
phi_start <- cur
theta_start <- cur + n_units
}
} else if (type == "lf") {
# Latent factor block stores u (n_latent) followed by lambda
# (n_arms). Total size = n_latent + n_arms.
block_size[b] <- n_units + n_arms
} else if (type == "hsgp_mo") {
# Multi-output HSGP stores K * M coefficients in output-major
# order: beta_{k, m} at offset k * M + m. K = n_arms.
block_size[b] <- n_arms * n_units
} else if (type == "mcar") {
# Separable MCAR stores p coupled fields field-major: field a at
# offset a * n_units within the block (size p * n_units).
block_size[b] <- as.integer(prepared[[b]]$n_fields) * n_units
field_starts <- c(field_starts, cur)
field_block_types <- c(field_block_types, "mcar")
field_n_fields <- c(field_n_fields, as.integer(prepared[[b]]$n_fields))
if (is.null(phi_start)) phi_start <- cur
} else if (type == "miid") {
# Multivariate IID stores p coupled per-group coefficient fields
# field-major: field a at offset a * n_groups within the block
# (size p * n_groups). A per-group RE, not a spatial field, so it
# does not register a field_start / phi_start alias.
block_size[b] <- as.integer(prepared[[b]]$n_fields) * n_units
} else {
block_size[b] <- n_units
if (type %in% c("icar", "car_proper")) {
field_starts <- c(field_starts, cur)
field_block_types <- c(field_block_types, type)
field_n_fields <- c(field_n_fields, 1L)
if (is.null(phi_start)) {
phi_start <- cur
}
}
}
cur <- cur + block_size[b]
}
n_x <- cur
out <- list(
n_arms = n_arms,
p = p_arm,
n_re = n_re,
beta_start = beta_start,
re_start = re_start,
block_start = block_start,
block_size = block_size,
n_x = n_x
)
if (!is.null(phi_start)) out$phi_start <- phi_start
if (!is.null(theta_start)) out$theta_start <- theta_start
if (length(field_starts) > 0L) {
out$field_starts <- field_starts
out$field_block_types <- field_block_types
out$field_n_fields <- field_n_fields
}
out
}
# Posterior moments for the multi-block joint grid. Mirrors the single-arm
# multi-block helper (`.nl_posterior_moments_multi`) and additionally
# attaches the weighted-quantile median + 2.5/97.5 CI for every joint-grid
# axis -- the calibrated headline summary for right-skewed scale-like
# hyperparameters (alpha at small n_pos, sigma/range/phi near boundary).
.joint_posterior_moments_multi <- function(res, prepared, axis_offsets,
joint_grid, cp,
int_weights = NULL,
support = .NL_SUPPORT_KINDS,
within = NULL) {
support <- match.arg(support)
within <- .nl_within_cell_mode(within)
w <- res$weights
# Joint moments across every column of joint_grid (including phi
# columns appended for per-arm dispersion overrides).
res$theta_mean <- as.numeric(crossprod(w, joint_grid))
names(res$theta_mean) <- colnames(joint_grid)
res$theta_sd <- sqrt(pmax(0, as.numeric(crossprod(w, joint_grid^2)) -
res$theta_mean^2))
names(res$theta_sd) <- colnames(joint_grid)
B <- length(prepared)
per_block_moments <- vector("list", B)
for (b in seq_len(B)) {
n_axes_b <- axis_offsets[b + 1L] - axis_offsets[b]
if (n_axes_b == 0L) {
# Block contributes no outer-grid axes (e.g. type = "lf"):
# empty moments, empty axis_cols.
per_block_moments[[b]] <- list(
type = tolower(prepared[[b]]$type),
mean = numeric(0),
sd = numeric(0),
axis_cols = integer(0)
)
next
}
cols <- seq.int(axis_offsets[b] + 1L, axis_offsets[b + 1L])
sub <- joint_grid[, cols, drop = FALSE]
block_mean <- as.numeric(crossprod(w, sub))
block_sd <- sqrt(pmax(0, as.numeric(crossprod(w, sub^2)) - block_mean^2))
# Strip the "b<N>." prefix on per-block moment names -- block index
# is implicit in the list position.
bare_names <- sub("^b[0-9]+\\.", "", colnames(sub))
names(block_mean) <- bare_names
names(block_sd) <- bare_names
per_block_moments[[b]] <- list(
type = tolower(prepared[[b]]$type),
mean = block_mean,
sd = block_sd,
axis_cols = cols
)
}
res$block_moments <- per_block_moments
# Each axis's SD from the estimator its own resolution calls for, the rule
# `.nl_posterior_moments()` applies on every other nested path.
res <- .nl_attach_axis_sd(res)
# Weighted-quantile median + 2.5/97.5 CI per axis. After the (sigma,
# alpha) reparameterization, alpha is a column of joint_grid like
# any other axis; this attaches the calibrated summary for every
# axis, not just alpha.
# Supplied whatever the support: a moment rule needs the domain to form its
# interval at all, and a density read needs it to place its outer cell edges
# inside the quantity's own support.
geo <- .joint_axis_geometry(list(theta_grid = joint_grid,
axis_offsets = axis_offsets,
blocks = prepared))
qs <- .nl_axis_quantiles(joint_grid, res$log_marginal,
res$refining_axis, weights = int_weights,
log_quad = if (is.null(int_weights)) res$log_quad,
support = support, domains = geo$domain,
within = within, atoms = geo$atom, sd = TRUE)
res$theta_median <- qs$median
res$theta_ci_lo <- qs$ci_lo
res$theta_ci_hi <- qs$ci_hi
res$within_cell_requested <- within
.nl_attach_interval_provenance(res, qs, joint_grid, geo$domain, geo$atom)
}
Try the tulpa package in your browser
Any scripts or data that you put into this service are public.
R Package Documentation
Browse R Packages
We want your feedback!
Note that we can't provide technical support on individual packages. You should contact the package authors for that.
Embedding an R snippet on your website
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.