tests/testthat/test-psis.R

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

# --------------------------------------------------------------------------- #
# Fast outer Pareto-k diagnostic helpers (gcol33/tulpa#118): near-neighbour     #
# batch ordering + the speed knobs. Pure-R, no fitting.                         #
# --------------------------------------------------------------------------- #

test_that(".joint_is_chain_order is a permutation seeded near the centre", {
  set.seed(11)
  X <- matrix(rnorm(200 * 3), 200, 3)
  center <- c(5, 5, 5)                                   # far corner
  ord <- tulpa:::.joint_is_chain_order(X, center)
  expect_setequal(ord, seq_len(nrow(X)))                 # a genuine permutation
  # First visited row is the one nearest `center` in standardised space.
  sds <- apply(X, 2L, sd); Z <- sweep(X, 2L, sds, `/`)
  expect_equal(ord[1], which.min(colSums((t(Z) - center / sds)^2)))
  # Consecutive standardised steps are short relative to a random order.
  step_len <- function(o) mean(sqrt(rowSums((Z[o[-1], ] - Z[o[-length(o)], ])^2)))
  expect_lt(step_len(ord), step_len(sample(nrow(X))))
  # Degenerate sizes pass through.
  expect_identical(tulpa:::.joint_is_chain_order(X[1:2, , drop = FALSE]), 1:2)
})

test_that(".joint_is_solve_reordered restores the caller's row order", {
  X <- matrix(rnorm(60 * 2), 60, 2)
  # solve_fn returns a per-row signature (row sum) so we can check un-permutation.
  solve_fn <- function(m) rowSums(m)
  out <- tulpa:::.joint_is_solve_reordered(X, center = c(0, 0), solve_fn)
  expect_equal(out, rowSums(X))                          # back in original order
  # A length mismatch (kernel failure) is passed through verbatim.
  bad <- tulpa:::.joint_is_solve_reordered(X, c(0, 0), function(m) numeric(0))
  expect_length(bad, 0L)
})

test_that(".nested_grid_radius_cap is the scaled max whitened grid radius", {
  # gcol33/tulpa#94: the cap is .K_DIAG_SUPPORT_RADIUS_MULT times the largest
  # whitened distance || L^{-1} (node - u_hat) || over the grid nodes.
  set.seed(101)
  d <- 3L
  u_hat <- c(0.5, -1, 2)
  A  <- matrix(rnorm(d * d), d)
  Su <- crossprod(A) + diag(d)                          # SPD
  L  <- t(chol(Su))
  # Place nodes at KNOWN whitened radii: node_k = u_hat + L %*% z_k.
  Zn <- rbind(c(1, 0, 0), c(0, 2, 0), c(0, 0, 1.5), c(1, 1, 1))
  u_grid <- t(apply(Zn, 1, function(z) u_hat + as.numeric(L %*% z)))
  cap <- .nested_grid_radius_cap(u_grid, u_hat, L)
  expect_equal(cap, max(sqrt(rowSums(Zn^2))) * .K_DIAG_SUPPORT_RADIUS_MULT,
               tolerance = 1e-10)
})

test_that(".nested_is_pareto_k declines without solving when radius_cap is too tight", {
  # When fewer than .PSIS_MIN_EVAL draws fall within support, the core returns
  # NA WITHOUT paying any inner solve (the floor check precedes log_target).
  d <- 2L; u_hat <- c(0, 0); L_scale <- diag(d)
  hit <- 0L
  tgt <- function(U) { hit <<- hit + nrow(U); rep(0, nrow(U)) }
  set.seed(404)
  kd <- .nested_is_pareto_k(u_hat, L_scale, tgt, n_samples = 1000L,
                            radius_cap = 0.01)            # ~0 draws survive
  expect_true(is.na(kd$pareto_k))
  expect_lt(kd$n_eval, .PSIS_MIN_EVAL)
  expect_identical(hit, 0L)                               # no inner solve paid
})

skip_on_cran_budget()
# Pareto-smoothed importance sampling: the native PSIS core, its agreement with
# loo, and the outer Pareto-k-hat diagnostic for the nested-Laplace path.

test_that("tulpa_psis reproduces loo::psis pareto_k", {
  skip_on_cran_budget()
  skip_if_not_installed("loo")
  set.seed(42)

  cases <- list(
    light = rnorm(3000),                 # lognormal weights: all moments finite
    heavy = log(1 / runif(3000)),        # Pareto(1) weights: k-hat ~ 1
    mixed = rnorm(2000) + rexp(2000)     # moderately heavy
  )
  for (nm in names(cases)) {
    lr  <- cases[[nm]]
    ours <- tulpa_psis(lr)$pareto_k
    ref  <- suppressWarnings(
      loo::psis(matrix(lr, ncol = 1), r_eff = NA)$diagnostics$pareto_k)
    expect_equal(ours, ref, tolerance = 0.02,
                 info = sprintf("case '%s': ours=%.4f loo=%.4f", nm, ours, ref))
  }
})

test_that("a tail tied at the cutoff is left unsmoothed, as loo::psis leaves it", {
  skip_on_cran_budget()
  skip_if_not_installed("loo")
  # A quarter of the exceedances exactly 0 takes the Zhang-Stephens scale anchor
  # to 0 and the fit to k = Inf with a NaN scale. loo answers the same k and
  # keeps the tail at its raw log ratios; the engine must not write NaN there.
  lr <- c(rep(0, 175), seq(0.5, 5, length.out = 25))
  ours <- tulpa_psis(lr)
  ref  <- suppressWarnings(
    loo::psis(matrix(lr, ncol = 1), r_eff = NA))

  expect_equal(ours$pareto_k,
               ref$diagnostics$pareto_k, tolerance = 1e-12)
  expect_false(ours$tail_smoothed)
  expect_false(any(is.nan(ours$log_weights)))
  expect_equal(ours$log_weights,
               as.numeric(loo::weights.importance_sampling(ref, log = TRUE)),
               tolerance = 1e-12)
})

test_that("tulpa_psis keeps one weight per draw on non-finite ratios (#892)", {
  skip_on_cran_budget()
  set.seed(1)
  lr <- rnorm(100)

  # +Inf is the worst ratio there is: an infinite shape, all the weight on it,
  # never a healthy k read off the draws left after dropping it.
  r <- tulpa_psis(c(lr, Inf))
  expect_identical(r$pareto_k, Inf)
  expect_length(r$log_weights, 101L)
  expect_equal(exp(r$log_weights), c(rep(0, 100), 1))
  expect_equal(r$is_ess, 1)

  # -Inf is a zero-weight draw that keeps its place.
  r <- tulpa_psis(c(lr[1:50], -Inf, lr[51:100]))
  expect_length(r$log_weights, 101L)
  expect_identical(r$log_weights[51], -Inf)
  expect_equal(sum(exp(r$log_weights)), 1)
  expect_true(is.finite(r$pareto_k))

  expect_error(tulpa_psis(c(1, 2, NA, 4)), "`log_ratios`")
  expect_error(tulpa_psis(c(lr, NaN)), "`log_ratios`")
  expect_error(tulpa_psis("a"), "`log_ratios`")
  expect_error(tulpa_psis(rep(-Inf, 10)), "all -Inf")

  # Too few draws for a tail fit: still one normalized weight per draw.
  r <- tulpa_psis(c(0, 1, 2))
  expect_true(is.na(r$pareto_k))
  expect_equal(r$log_weights, c(0, 1, 2) - log(sum(exp(0:2))))
})

test_that("a -Inf ratio is the limit of a vanishing finite one (#892)", {
  skip_on_cran_budget()
  # A zero-weight draw is still a draw from the proposal: it sits below the
  # tail, counts towards S, and weighs exactly what a ratio far enough down
  # that its exponential underflows weighs.
  set.seed(2)
  lr <- c(rnorm(300) + rexp(300), rep(-Inf, 20))[sample(320)]
  lr_fin <- lr
  lr_fin[!is.finite(lr)] <- min(lr[is.finite(lr)]) - 1e4
  ours <- tulpa_psis(lr)
  ref  <- tulpa_psis(lr_fin)
  expect_equal(ours$pareto_k, ref$pareto_k, tolerance = 1e-12)
  expect_equal(ours$tail_len, ref$tail_len)
  expect_equal(exp(ours$log_weights), exp(ref$log_weights), tolerance = 1e-12)
  expect_true(all(ours$log_weights[!is.finite(lr)] == -Inf))
})

test_that(".kdiag_knobs defaults to the fast path and honours overrides", {
  skip_on_cran_budget()
  d <- tulpa:::.kdiag_knobs()
  expect_identical(d$refresh, tulpa:::.K_DIAG_REFRESH)
  expect_identical(d$tol,     tulpa:::.K_DIAG_TOL)
  expect_true(d$reorder)
  expect_true(d$percell)
  withr::local_options(tulpa.kdiag.refresh = 1L, tulpa.kdiag.tol = 0,
                       tulpa.kdiag.reorder = FALSE, tulpa.kdiag.percell = FALSE)
  e <- tulpa:::.kdiag_knobs()
  expect_identical(e$refresh, 1L)
  expect_identical(e$tol, 0)
  expect_false(e$reorder)
  expect_false(e$percell)
})

test_that(".joint_nearest_grid_mode picks each draw's nearest grid cell mode", {
  skip_on_cran_budget()
  # 3 grid cells at theta = 1, 5, 9 (1-D), with modes a fixed function of the
  # cell so the picked row is identifiable.
  res <- list(theta_grid = matrix(c(1, 5, 9), ncol = 1),
              modes = rbind(c(10, 10), c(50, 50), c(90, 90)))
  draws <- matrix(c(0.8, 4.6, 9.3, 5.2), ncol = 1)   # nearest: 1, 5, 9, 5
  W <- tulpa:::.joint_nearest_grid_mode(draws, res)
  expect_equal(dim(W), c(4L, 2L))
  expect_equal(W[, 1], c(10, 50, 90, 50))
  # Declines (NULL) when modes are missing or mis-shaped -> caller broadcasts.
  expect_null(tulpa:::.joint_nearest_grid_mode(draws, list(theta_grid = res$theta_grid)))
  expect_false(tulpa:::.joint_has_grid_modes(list(theta_grid = res$theta_grid)))
  expect_true(tulpa:::.joint_has_grid_modes(res))
})

test_that(".joint_make_diag_refit selects per-cell warm when modes are present", {
  skip_on_cran_budget()
  res <- list(theta_grid = matrix(c(1, 5, 9), ncol = 1),
              modes = rbind(c(10, 10), c(50, 50), c(90, 90)))
  seen <- NULL
  solve_fn <- function(theta_mat, x_init_per_cell = NULL) {
    seen <<- x_init_per_cell           # capture what the strategy passed down
    rowSums(theta_mat)
  }
  knobs <- list(refresh = 4L, tol = 1e-4, reorder = TRUE, percell = TRUE)
  refit <- tulpa:::.joint_make_diag_refit(res, solve_fn, c(5), knobs)
  out <- refit(matrix(c(0.8, 4.6, 9.3), ncol = 1))
  expect_equal(out, c(0.8, 4.6, 9.3))   # row order preserved (no re-order)
  expect_equal(dim(seen), c(3L, 2L))    # per-cell warm matrix was passed
  # No modes -> falls back to the re-order strategy (no per-cell matrix).
  seen <- NULL
  refit2 <- tulpa:::.joint_make_diag_refit(list(theta_grid = res$theta_grid),
                                           solve_fn, c(5), knobs)
  invisible(refit2(matrix(c(0.8, 4.6, 9.3, 5.0), ncol = 1)))
  expect_null(seen)
})

test_that("pareto_k separates light from heavy tails; is_ess is bounded", {
  skip_on_cran_budget()
  set.seed(1)
  light <- tulpa_psis(rnorm(3000))
  heavy <- tulpa_psis(log(1 / runif(3000)))

  expect_lt(light$pareto_k, 0.5)         # finite-moment proposal: reliable
  expect_gt(heavy$pareto_k, 0.7)         # Pareto(1): not correctable
  expect_lt(light$is_ess, 3000 + 1e-6)   # ESS never exceeds the sample size
  expect_gt(light$is_ess, 0)
  expect_lt(heavy$is_ess, light$is_ess)  # heavy weights concentrate -> lower ESS
})

test_that(".nested_outer_pareto_k is low when the proposal covers the target", {
  skip_on_cran_budget()
  set.seed(7)
  # Gaussian target N(mu, Sigma); proposal slightly wider (1.5 * Sigma) so the
  # importance weights are bounded -> small / negative k-hat, high IS-ESS.
  mu    <- c(0.4, -0.2)
  Sigma <- matrix(c(1, 0.3, 0.3, 0.5), 2)
  Sinv  <- solve(Sigma)
  log_target <- function(U) {
    D <- sweep(U, 2L, mu)
    -0.5 * rowSums((D %*% Sinv) * D)                  # up to a constant
  }
  L_scale <- t(chol(1.5 * Sigma))                     # proposal scale > target

  kd <- .nested_outer_pareto_k(log_target, theta_hat = mu,
                               L_scale = L_scale, n_samples = 2000L)
  expect_true(is.finite(kd$pareto_k))
  expect_lt(kd$pareto_k, 0.5)
  expect_gt(kd$is_ess, 0)
  expect_equal(kd$n_eval, 2000L)
})

test_that(".nested_outer_pareto_k rises when the target is heavier than the proposal", {
  skip_on_cran_budget()
  set.seed(9)
  mu <- c(0, 0)
  # Student-t(df = 2) target (infinite variance) against a unit-Gaussian
  # proposal of matching scale: tails the Gaussian cannot cover -> high k-hat.
  log_target_t <- function(U) {
    rowSums(stats::dt(sweep(U, 2L, mu), df = 2, log = TRUE))
  }
  kd <- .nested_outer_pareto_k(log_target_t, theta_hat = mu,
                               L_scale = diag(1, 2), n_samples = 4000L)
  expect_gt(kd$pareto_k, 0.5)                          # heavier target -> not correctable
})

test_that("the IS cores decline before evaluating the target below the floor", {
  skip_on_cran_budget()
  # gcol33/tulpa#51: a sub-floor n_samples can never reach .PSIS_MIN_EVAL finite
  # evaluations, so both cores must return NA WITHOUT touching the (expensive)
  # target -- not evaluate the whole budget and then discard it.
  expect_lt(20L, .PSIS_MIN_EVAL)

  hit <- 0L
  target_counting <- function(U) { hit <<- hit + nrow(U); rep(0, nrow(U)) }
  kd <- .nested_outer_pareto_k(target_counting, theta_hat = c(0, 0),
                               L_scale = diag(1, 2), n_samples = 20L)
  expect_true(is.na(kd$pareto_k))
  expect_equal(kd$n_eval, 0L)
  expect_identical(hit, 0L)                            # no inner solve was paid

  hit_b <- 0L
  batched_counting <- function(U) { hit_b <<- hit_b + nrow(U); rep(0, nrow(U)) }
  kb <- .nested_is_pareto_k(theta_hat = c(0, 0), L_scale = diag(1, 2),
                            log_target_batched = batched_counting, n_samples = 20L)
  expect_true(is.na(kb$pareto_k))
  expect_equal(kb$n_eval, 0L)
  expect_identical(hit_b, 0L)
})

test_that(".nested_is_pareto_k skips draws beyond radius_cap (cost bound, light tail)", {
  skip_on_cran_budget()
  # gcol33/tulpa#94: for a target NO heavier than the Gaussian proposal a finite
  # radius_cap excludes deep-extrapolation draws WITHOUT an inner solve, bounding
  # the diagnostic cost; radius_cap = Inf reproduces the unrestricted path. The
  # target here is narrower than the proposal, so the importance log-ratio FALLS
  # with the radius and the escalation step (gcol33/tulpa#100) does not fire.
  d <- 2L; u_hat <- c(0, 0); L_scale <- diag(d)
  rows_seen <- 0L
  tgt <- function(U) { rows_seen <<- rows_seen + nrow(U); -0.75 * rowSums(U^2) }

  # Z ~ N(0, I_2): ||z||^2 ~ chi^2_2. Cap radius 2 keeps z2 <= 4, i.e. a
  # fraction pchisq(4, 2) = 1 - exp(-2) ~ 0.865 of the draws.
  set.seed(303)
  kd <- .nested_is_pareto_k(u_hat, L_scale, tgt, n_samples = 2000L,
                            radius_cap = 2)
  expect_lt(rows_seen, 2000L)                  # the deep tail was skipped
  expect_gt(rows_seen, 0.70 * 2000L)           # the bulk was retained
  expect_equal(kd$n_eval, rows_seen)           # every retained draw had finite target

  rows_seen <- 0L
  set.seed(303)
  kd_inf <- .nested_is_pareto_k(u_hat, L_scale, tgt, n_samples = 2000L,
                                radius_cap = Inf)
  expect_equal(rows_seen, 2000L)               # Inf cap == evaluate every draw
})

test_that(".nested_is_pareto_k folds the clipped tail back in for a heavy target (#100)", {
  skip_on_cran_budget()
  # gcol33/tulpa#100: the radius cap must NOT bias k-hat downward when the target
  # is genuinely heavier-tailed than the Gaussian proposal -- the regime a high
  # k-hat exists to flag. There the far-radius draws carry the LARGEST importance
  # ratios, so the escalation step re-evaluates them and the capped k-hat tracks
  # the uncapped one instead of collapsing toward "reliable".
  d <- 1L; u_hat <- 0; L_scale <- matrix(1, 1, 1)
  heavy <- function(U) stats::dt(U[, 1], df = 3, log = TRUE)   # heavier than N(0,1)

  k_at <- function(cap, seed = 17) {
    set.seed(seed)
    .nested_is_pareto_k(u_hat, L_scale, heavy, n_samples = 6000L,
                        radius_cap = cap)$pareto_k
  }
  k_inf <- k_at(Inf)
  expect_true(is.finite(k_inf))
  expect_gt(k_inf, 0.5)                         # the t_3 / Gaussian ratio is heavy
  for (cap in c(2.2, 2.5, 3.0)) {               # realistic >= 2x-grid-coverage caps
    k_cap <- k_at(cap)
    expect_true(is.finite(k_cap))
    # Not materially below the uncapped reading: the escalation prevents the
    # downward bias the cap would otherwise introduce.
    expect_gt(k_cap, k_inf - 0.12)
  }

  # Control: a Gaussian target matched to the proposal stays light under the cap
  # (no escalation, the cost bound is preserved).
  light <- function(U) stats::dnorm(U[, 1], log = TRUE)
  set.seed(17)
  k_light <- .nested_is_pareto_k(u_hat, L_scale, light, n_samples = 6000L,
                                 radius_cap = 2.2)$pareto_k
  expect_true(is.na(k_light) || k_light < 0.7)
})

skip_on_cran()

test_that("tulpa_re_cov_nested reports a Pareto-k-hat without disturbing draws", {
  skip_on_cran()
  sim <- function(seed, G = 50L, npg = 10L) {
    set.seed(seed)
    N <- G * npg; grp <- rep(seq_len(G), each = npg)
    x <- rnorm(N); X <- cbind(1, x); Z <- cbind(1, x)
    Sig <- matrix(c(0.64, 0.24, 0.24, 0.36), 2)
    u <- t(t(chol(Sig)) %*% matrix(rnorm(2 * G), 2))
    eta <- as.numeric(X %*% c(-0.3, 0.7)) + rowSums(Z * u[grp, ])
    list(y = rbinom(N, 1L, plogis(eta)), X = X, Z = Z, grp = grp, G = G, N = N)
  }
  d  <- sim(11L)
  rt <- list(idx = d$grp, n_groups = d$G, n_coefs = 2L, Z = d$Z)

  fit <- tulpa_re_cov_nested(d$y, rep(1L, d$N), d$X, rt, family = "binomial",
                             control = list(seed = 11L, k_samples = 200L))
  # k-hat is a finite, data-dependent reading (a sparse binary RE-covariance
  # posterior is skewed, so a HIGH k-hat here is a correct signal, not a bug --
  # the estimator's correctness is pinned by the loo-equivalence test above and
  # the synthetic-target helper test). Assert the plumbing and ESS range only.
  expect_true(is.finite(fit$pareto_k))
  expect_true(is.finite(fit$pareto_k_is_ess))
  expect_gt(fit$pareto_k_is_ess, 0)
  expect_lte(fit$pareto_k_is_ess, 200 + 1e-6)
  expect_equal(fit$pareto_k_scope, "outer (hyperparameter) Gaussian proposal")
  # A computed k-hat carries no decline reason (gcol33/tulpa#295).
  expect_true(is.na(fit$pareto_k_declined))

  # diagnose_k must not perturb the fixed-effect draws (RNG state restored).
  off <- tulpa_re_cov_nested(d$y, rep(1L, d$N), d$X, rt, family = "binomial",
                             control = list(seed = 11L, diagnose_k = FALSE))
  on  <- tulpa_re_cov_nested(d$y, rep(1L, d$N), d$X, rt, family = "binomial",
                             control = list(seed = 11L, diagnose_k = TRUE,
                                            k_samples = 150L))
  expect_true(is.na(off$pareto_k))
  # The NA says WHY it is NA: nothing is wrong, the caller asked for it off.
  expect_identical(off$pareto_k_declined, "not_requested")
  expect_equal(off$draws, on$draws)
})

test_that("tulpa_nested_laplace reports an outer k-hat for a positive-scale block", {
  skip_on_cran()
  set.seed(3)
  nr <- 60L; spr <- 10L; N <- nr * spr
  region <- rep(seq_len(nr), each = spr)
  X <- cbind(1, rnorm(N))
  u <- rnorm(nr, 0, 0.5)
  y <- rbinom(N, 10L, plogis(as.numeric(X %*% c(-0.2, 0.7)) + u[region]))
  sg <- exp(seq(log(0.2), log(1.5), length.out = 7))
  mk <- function(grid) list(list(type = "iid", obs_idx = region,
                                 n_units = nr, sigma_grid = grid))

  fit <- suppressWarnings(tulpa_nested_laplace(
    y = y, n_trials = rep(10L, N), X = X, prior = mk(sg),
    family = "binomial", phi = 1,
    control = list(max_iter = 100L, tol = 1e-8, k_samples = 200L)))
  # Single positive-scale (RE-SD) axis: k-hat is computed via the log transform.
  # The value is data-dependent (a sparse binary RE-SD posterior is skewed, so a
  # high reading is correct) -- assert plumbing + ESS range, as in the re_cov case.
  expect_true(is.finite(fit$pareto_k))
  expect_true(is.finite(fit$pareto_k_is_ess))
  expect_gt(fit$pareto_k_is_ess, 0)
  expect_lte(fit$pareto_k_is_ess, 200 + 1e-6)
  expect_equal(fit$pareto_k_scope, "outer (hyperparameter) Gaussian proposal")

  off <- suppressWarnings(tulpa_nested_laplace(
    y = y, n_trials = rep(10L, N), X = X, prior = mk(sg),
    family = "binomial", phi = 1,
    control = list(max_iter = 100L, tol = 1e-8, diagnose_k = FALSE)))
  expect_true(is.na(off$pareto_k))                  # gated off
})

test_that("outer k-hat declines (NA) for a multi-block nested fit", {
  skip_on_cran()
  set.seed(4)
  nr <- 40L; spr <- 8L; N <- nr * spr
  region  <- rep(seq_len(nr), each = spr)
  region2 <- rep(seq_len(nr), times = spr)
  X <- cbind(1, rnorm(N))
  y <- rbinom(N, 5L, plogis(as.numeric(X %*% c(-0.1, 0.5))))
  sg <- exp(seq(log(0.2), log(1.2), length.out = 5))
  prior <- list(
    list(type = "iid", obs_idx = region,  n_units = nr, sigma_grid = sg),
    list(type = "iid", obs_idx = region2, n_units = nr, sigma_grid = sg))
  fit <- suppressWarnings(tulpa_nested_laplace(
    y = y, n_trials = rep(5L, N), X = X, prior = prior,
    family = "binomial", phi = 1, control = list(max_iter = 80L, tol = 1e-7)))
  expect_true(is.na(fit$pareto_k))                  # multi-block: declined, not guessed
  expect_gt(length(fit$weights), 1L)                # quadrature-ESS fallback available
})

# --------------------------------------------------------------------------- #
# Discriminating power on REAL engine output (gcol33/tulpa#99).                #
#                                                                              #
# The synthetic-target tests above show the estimator separates light from    #
# heavy tails; the real-fit tests above assert only plumbing, with a comment  #
# that a high k-hat "is a correct signal". This converts that comment into a  #
# tested invariant: the outer k-hat must ORDER two real RE-covariance fits --  #
# a well-identified one (many large groups -> near-Gaussian Sigma posterior,   #
# correctable) below a tiny-binary one (few binary obs per group -> strongly   #
# skewed variance posterior, not correctable). A regression that made the      #
# real-fit k-hat constant (always ~0, or always NA) would pass every other     #
# real-fit test but fail this one.                                             #
# --------------------------------------------------------------------------- #

test_that("a tiny-binary RE-covariance k-hat is a proposal scale, not a skewed posterior", {
  skip_on_cran()

  # hyperprior = "proper" stated explicitly: this test calibrates the k-hat
  # SEPARATION between a well-identified and a pathological fit, which is the PC
  # prior's shrinkage behavior at small G, not a flat-in-log prior's.
  #
  # This test used to assert that the two regimes separate by orders of
  # magnitude, on the reading that a small-group binary variance-component
  # posterior is genuinely skewed and its huge k-hat is therefore a correct
  # signal. Measured per candidate layer (gcol33/tulpa#630,
  # dev_notes/issue630/RESULTS630.md), it was not: the separation was the
  # PROPOSAL, mis-scaled on the tiny side, and it disappears once the proposal
  # is re-estimated from its own PSIS-weighted moments -- a candidate the joint
  # path has always scored and this backend did not. So what is asserted here is
  # the repair and its two limits, which is what a regression would undo.
  #
  # (a) Well-identified: 30 groups x 25 gaussian obs each, correlated random
  # slope. The Sigma posterior is near-Gaussian, so the proposal as PLACED
  # already fits and refinement has little to do.
  k_well <- function(seed) {
    set.seed(seed); G <- 30L; per <- 25L; n <- G * per
    grp <- rep(seq_len(G), each = per); x <- rnorm(n)
    eta <- 0.2 + 0.5 * x + rnorm(G, 0, 0.8)[grp] + rnorm(G, 0, 0.5)[grp] * x
    y   <- eta + rnorm(n, 0, 0.5)
    rt  <- list(idx = grp, n_groups = G, n_coefs = 2L, Z = cbind(1, x),
                correlated = TRUE)
    tulpa_re_cov_nested(y, rep(1L, n), cbind(1, x), rt, family = "gaussian",
                        phi = 0.25, hyperprior = "proper",
                        control = list(diagnose_k = TRUE, k_samples = 150L))
  }
  # (b) Tiny binary groups: 25 groups x 3 binary obs each. The proposal as
  # placed is badly scaled there, so its FIRST-PASS k-hat blows up.
  k_tiny <- function(seed) {
    set.seed(seed); G <- 25L; per <- 3L; n <- G * per
    grp <- rep(seq_len(G), each = per); x <- rnorm(n)
    eta <- -0.2 + 0.4 * x + rnorm(G, 0, 1.0)[grp] + rnorm(G, 0, 0.8)[grp] * x
    y   <- rbinom(n, 1L, plogis(eta))
    rt  <- list(idx = grp, n_groups = G, n_coefs = 2L, Z = cbind(1, x),
                correlated = TRUE)
    tulpa_re_cov_nested(y, rep(1L, n), cbind(1, x), rt, family = "binomial",
                        hyperprior = "proper",
                        control = list(diagnose_k = TRUE, k_samples = 150L))
  }

  fa <- lapply(1:5, function(s) k_well(100L + s))
  fb <- lapply(1:5, function(s) k_tiny(200L + s))
  get <- function(fits, nm) vapply(fits, function(f) f[[nm]], numeric(1))
  ka  <- get(fa, "pareto_k");            kb  <- get(fb, "pareto_k")
  fpa <- get(fa, "pareto_k_first_pass"); fpb <- get(fb, "pareto_k_first_pass")

  # 1. The first pass reproduces the old reading: the placement IS bad on the
  #    tiny side and fine on the well-identified one. That half was real.
  expect_true(all(is.finite(fpa)))
  expect_lt(median(fpa), 2.0)
  expect_gt(median(fpb), 5 * median(fpa))

  # 2. Refinement repairs it. Every well-identified fit stays finite, and the
  #    tiny-binary median falls by more than an order of magnitude off its own
  #    first pass -- so the huge number was the proposal's scale, not the
  #    posterior's shape.
  expect_true(all(is.finite(ka)))
  expect_lt(median(ka), 2.0)
  expect_lt(median(kb), median(fpb) / 10)

  # 3. The report says which proposal each k-hat belongs to. Every tiny-binary
  #    placement is past the escalation threshold; a fit whose k-hat comes from
  #    the moment-matched proposal improved on its first pass, and one left on
  #    the mode-Hessian proposal reports that first pass unchanged. Whether a
  #    given seed's rescue succeeds is not asserted: these modes sit on a ridge
  #    where the slope variance collapses (log sigma_2 near -11), and a
  #    rounding-level change in the inner solve moves the stopping point along
  #    it by ~0.4, enough to turn seed 203's rescue from failing (k 68) to
  #    succeeding (k 0.66) with its per-draw log-marginals unchanged to 2e-12.
  expect_true(all(fpb > .nl_diag("k_usable")))
  src <- vapply(c(fa, fb), function(f) f$pareto_k_proposal_source, "")
  expect_true(all(src %in% c("mode_hessian", "moment_matched")))
  k_all  <- c(ka, kb)
  fp_all <- c(fpa, fpb)
  mm <- src == "moment_matched"
  expect_true(all(k_all[mm] < fp_all[mm]))
  expect_identical(k_all[!mm], fp_all[!mm])
})