tests/testthat/test-loo-loco.R

In INLAvaan: Approximate Bayesian Latent Variable Analysis

# Extended LOO suite pinned to reference values. It runs in CI, and
# test-loo-loso.R covers the core LOO on CRAN.
skip_on_cran()

twolevel_model <- "
  level: 1
    fw =~ y1 + y2 + y3
    fw ~ x1 + x2 + x3
  level: 2
    fb =~ y1 + y2 + y3
    fb ~ w1 + w2
"

# Shrunk two-level fixture: first 24 cluster ids (6 full cycles of the
# 5/10/15/20 cluster-size pattern that repeats every 4 cluster ids in
# lavaan::Demo.twolevel), 300 rows total. Shared across both the main
# `fit` below and `fit_fx` (the fixed.x = TRUE variant further down).
d_sub <- lavaan::Demo.twolevel[lavaan::Demo.twolevel$cluster %in% 1:24, ]

fit <- asem(
  twolevel_model,
  d_sub,
  cluster = "cluster",
  meanstructure = TRUE,
  fixed.x = FALSE,
  verbose = FALSE,
  nsamp = 3,
  test = "none",
  vb_correction = FALSE,
  marginal_method = "marggaus",
  marginal_correction = "none"
)
res <- loo(fit)

test_that("LOCO matches reference values", {
  # Reference values computed with an independent implementation of the same
  # Taylor LOO formulas on this exact fit
  expect_equal(res$type, "loco")
  expect_equal(res$n_units, 24L)
  expect_equal(res$elpd_1, -2811.3389254026, tolerance = 1e-4)
  expect_equal(res$elpd_2, -2832.9728115262, tolerance = 1e-4)
  expect_equal(res$se_1, 259.1529664233, tolerance = 1e-4)
  expect_equal(res$se_2, 260.9893088931, tolerance = 1e-4)
  # p_loo aggregates small differences of nearly-equal numbers, so it is
  # more sensitive to BLAS/optimiser endpoint noise than the elpd totals
  expect_equal(res$p_loo_1, 30.7792227732, tolerance = 1e-2)
  expect_equal(res$p_loo_2, 35.1657403910, tolerance = 1e-2)

  # cluster ids 1, 6, 8 have sizes 5, 10, 20 respectively under the
  # 5/10/15/20 cycle
  pu <- res$per_unit[c(1L, 6L, 8L), ]
  expect_equal(pu$nobs, c(5L, 10L, 20L))
  expect_equal(
    pu$l_star,
    c(-45.1868700825, -91.3569642108, -177.6891056793),
    tolerance = 1e-4
  )
  expect_equal(
    pu$log_cpo_1,
    c(-45.3396350528, -91.5858133174, -178.4008589303),
    tolerance = 1e-4
  )
  expect_equal(
    pu$log_cpo_2,
    c(-45.5376421501, -92.0469159588, -179.5548931984),
    tolerance = 1e-4
  )
  # det_term values are tiny finite-difference remainders whose relative
  # error is dominated by cross-platform noise; their correctness is pinned
  # through log_cpo_2 above, so only check structure here
  expect_true(all(is.finite(pu$det_term)))
  expect_true(all(pu$det_term < 0))
})

test_that("LOCO structure and internal identities", {
  expect_s3_class(res, "inlavaan_loo")
  expect_true(all(res$per_unit$ok))
  expect_equal(sum(res$per_unit$nobs), nrow(d_sub))
  expect_equal(
    res$per_unit$lpd_1 + res$per_unit$log_cpo_1,
    2 * res$per_unit$l_star
  )
  expect_output(print(res), "Leave-one-cluster-out")
})

test_that("the curvature check is reported and printed", {
  # The check compares the summed first-to-second-order gap against pD/2, both
  # read off the same Laplace summary
  expect_equal(res$pd_trace, sum(res$per_unit$k_sum))
  expect_equal(res$elpd_gap, res$elpd_1 - res$elpd_2)

  # Pin the width so the cli rules and the reflowed note render the same way
  # whatever console the tests run on
  old_opt <- options(cli.width = 100)
  on.exit(options(old_opt), add = TRUE)

  expect_output(print(res), "Curvature check")
  expect_output(print(res), "pD/2 \\(trace\\)")
  expect_output(print(res), "excess over pD/2")
  # summary() is an alias for print(), not a different view
  expect_identical(capture.output(print(res)), capture.output(summary(res)))

  # At first order there is no second-order score to take a gap against, so
  # both scalars are NA and the block is skipped entirely
  res_fo <- loo(fit, second_order = FALSE)
  expect_true(is.na(res_fo$pd_trace))
  expect_true(is.na(res_fo$elpd_gap))
  expect_false(any(grepl("Curvature check", capture.output(print(res_fo)))))

  # A units subset sums both sides over the same units, so the check still
  # holds while the totals are partial
  res_sub <- loo(fit, units = 1:2)
  expect_equal(res_sub$pd_trace, sum(res_sub$per_unit$k_sum))
  expect_lt(res_sub$pd_trace, res$pd_trace)

  # Sum of cluster logliks equals the model loglik at the mode
  int <- get_inlavaan_internal(fit)
  x <- INLAvaan:::pars_to_x(int$theta_star, int$partable)
  lm_x <- lavaan::lav_model_set_parameters(int$lavmodel, x)
  opts <- fit@Options
  opts$estimator <- "ML"
  ll <- lavaan:::lav_model_loglik(
    lavdata = int$lavdata,
    lavsamplestats = int$lavsamplestats,
    lavimplied = lavaan::lav_model_implied(lm_x),
    lavmodel = lm_x,
    lavoptions = opts
  )$loglik
  expect_equal(sum(res$per_unit$l_star), ll, tolerance = 1e-6)
})

test_that("curvature diagnostics satisfy their defining identities", {
  pu <- res$per_unit
  # k_u = lambda_max(-Omega H_u) >= 0 when the unit curvature is n.s.d.
  expect_true(all(pu$k_max >= 0))
  # Finite mean of the importance ratio <=> A_u positive definite <=> ok
  expect_identical(pu$ok, pu$k_max < 1)
  # det_term = 1/2 sum_i log(1 - k_i) <= 1/2 log(1 - k_max), so k_max is
  # bounded by the determinant term alone
  expect_true(all(pu$k_max <= 1 - exp(2 * pu$det_term) + 1e-10))
  # Trace dominates the largest eigenvalue
  expect_true(all(pu$k_sum >= pu$k_max - 1e-10))
  # First-order scoring computes no curvature, so both are NA
  res1 <- loo(fit, second_order = FALSE)
  expect_true(all(is.na(res1$per_unit$k_max)))
  expect_true(all(is.na(res1$per_unit$k_sum)))
})

test_that("a missing second-order term takes the whole statistic to first order", {
  S <- get_inlavaan_internal(fit)$Sigma_theta
  # Inflating Omega scales every k_u linearly, driving A_u = Omega^-1 + H_u
  # toward the negative-definite H_u, so units lose their second-order term
  # deterministically rather than by a lucky data seed.

  # Nothing fails here: the second-order total is the plain sum, no warning.
  expect_equal(res$n_ok, res$n_units)
  expect_true(res$use_second)
  expect_equal(res$elpd_2, sum(res$per_unit$log_cpo_2))
  expect_no_warning(loo(fit, cores = 1L))

  # Every unit fails: there is no second-order total to report at all.
  r_all <- suppressWarnings(loo(fit, Omega = S * 1e6, cores = 1L))
  expect_equal(r_all$n_ok, 0L)
  expect_true(all(is.na(r_all$per_unit$log_cpo_2)))
  expect_true(is.na(r_all$elpd_2))
  expect_false(r_all$use_second)
  expect_equal(unname(r_all$estimates["elpd_loo", "Estimate"]), r_all$elpd_1)

  # Some units fail, and that is enough: elpd_loo is the first-order total
  # over all n units. It is neither a blend of the two orders nor the
  # second-order sum over the surviving units, which would score the model
  # over fewer units and so flatter it.
  r_mix <- suppressWarnings(loo(fit, Omega = S * 4, cores = 1L))
  pu <- r_mix$per_unit
  expect_true(any(is.na(pu$log_cpo_2)))
  expect_false(all(is.na(pu$log_cpo_2)))
  expect_true(is.na(r_mix$elpd_2))
  expect_equal(unname(r_mix$estimates["elpd_loo", "Estimate"]), r_mix$elpd_1)
  blended <- ifelse(is.na(pu$log_cpo_2), pu$log_cpo_1, pu$log_cpo_2)
  expect_false(isTRUE(all.equal(r_mix$elpd_1, sum(blended))))
  expect_false(isTRUE(all.equal(r_mix$elpd_1, sum(pu$log_cpo_2, na.rm = TRUE))))

  # The pointwise column records the same failures the aggregates react to.
  expect_true(all(is.na(pu$log_cpo_2) == !pu$ok))

  # A missing log CPO term warns; it is the condition that says leave-one-out
  # is not identified for that unit
  msg <- tryCatch(
    loo(fit, Omega = S * 4, cores = 1L),
    warning = conditionMessage
  )
  expect_match(msg, "no second-order term")
  # ... and it names them, since Omega^-1 + H_u is the deleted posterior
  # precision, so the finding is about the unit and the user's next move is
  # to go and look at it
  bad <- pu$unit[!pu$ok]
  expect_true(all(vapply(
    bad,
    function(u) grepl(paste0("\\b", u, "\\b"), msg),
    logical(1)
  )))
  expect_match(msg, "per_unit", fixed = TRUE)
})

test_that("LOCO unit subsetting and theta/Omega override", {
  sub <- c(3L, 7L, 11L)
  res_sub <- loo(fit, units = sub)
  expect_equal(res_sub$per_unit$unit, sub)
  expect_equal(
    res_sub$per_unit$log_cpo_2,
    res$per_unit$log_cpo_2[sub],
    tolerance = 1e-8
  )

  int <- get_inlavaan_internal(fit)
  res_same <- loo(
    fit,
    theta = int$theta_star,
    Omega = int$Sigma_theta,
    units = sub
  )
  expect_true(res_same$theta_overridden)
  expect_equal(
    res_same$per_unit$log_cpo_2,
    res$per_unit$log_cpo_2[sub],
    tolerance = 1e-8
  )
})

test_that("two-level LOSO override scores row deletions", {
  expect_warning(
    res_row <- loo(fit, type = "loso", units = 1:8),
    "leave-one-unit-out"
  )
  expect_equal(res_row$type, "loso")
  expect_equal(nrow(res_row$per_unit), 8L)
  expect_true(all(res_row$per_unit$ok))
  expect_true(all(res_row$per_unit$nobs == 1L))

  int <- get_inlavaan_internal(fit)
  css <- INLAvaan:::loco_suff_stats(int$lavdata)
  X <- int$lavdata@X[[1L]]
  cache <- INLAvaan:::loo_grad_cache(
    int$theta_star,
    int$lavmodel,
    int$partable,
    two_level = TRUE
  )

  # The downdated sufficient statistics agree with rebuilding the
  # cluster-minus-row statistics from the raw data
  i <- 2L
  j <- css$cluster_idx[i]
  rows <- setdiff(which(css$cluster_idx == j), i)
  Yj <- X[rows, , drop = FALSE]
  us_raw <- INLAvaan:::loco_stats_build(
    length(rows),
    crossprod(sweep(Yj, 2L, colMeans(Yj), "-")),
    colMeans(Yj)[css$zy_idx],
    css
  )
  ll_minus_raw <- INLAvaan:::loco_loglik_us(us_raw, cache$mom)
  ll_full <- INLAvaan:::loco_loglik_one(j, css, cache$mom)
  expect_equal(
    res_row$per_unit$l_star[i],
    ll_full - ll_minus_raw,
    tolerance = 1e-8
  )

  # Analytic row score matches a numerical derivative
  s2 <- as.numeric(INLAvaan:::loso2l_scores_theta(
    int$theta_star,
    css,
    X,
    int$lavmodel,
    int$partable,
    units = i
  ))
  h <- 1e-6
  g_num <- vapply(
    seq_along(int$theta_star),
    function(k) {
      tp <- tm <- int$theta_star
      tp[k] <- tp[k] + h
      tm[k] <- tm[k] - h
      cp <- INLAvaan:::loo_grad_cache(
        tp,
        int$lavmodel,
        int$partable,
        two_level = TRUE
      )
      cm <- INLAvaan:::loo_grad_cache(
        tm,
        int$lavmodel,
        int$partable,
        two_level = TRUE
      )
      (INLAvaan:::loso2l_loglik_all(i, css, X, cp$mom) -
        INLAvaan:::loso2l_loglik_all(i, css, X, cm$mom)) /
        (2 * h)
    },
    numeric(1)
  )
  expect_equal(s2, g_num, tolerance = 1e-5)
})

test_that("fixed.x two-level fits are scored conditionally", {
  fit_fx <- asem(
    twolevel_model,
    d_sub,
    cluster = "cluster",
    meanstructure = TRUE,
    fixed.x = TRUE,
    verbose = FALSE,
    nsamp = 3,
    test = "none",
    vb_correction = FALSE,
    marginal_method = "marggaus",
    marginal_correction = "none"
  )
  res_fx <- loo(fit_fx, units = 1:5)
  expect_equal(res_fx$flavour, "conditional")
  expect_true(all(is.finite(res_fx$per_unit$log_cpo_2)))
})

test_that("waic gains type: conditional (leave-one-unit-out) WAIC", {
  # default is marginal (per-cluster) WAIC
  w_marg <- suppressWarnings(waic(fit))
  expect_equal(w_marg$type, "loco")

  # type = "loso" warns and scores the conditional (leave-one-unit-out)
  # WAIC, the same estimand as loo(type = "loso"): identical lpd terms, so
  # the two differ only by the p_waic vs p_loo penalty gap
  w_cond <- testthat::capture_warnings(
    w <- waic(fit, type = "loso", units = 1:5)
  )
  expect_true(any(grepl("leave-one-unit-out", w_cond)))
  expect_equal(w$type, "loso")
  expect_equal(w$n_units, 5L)
  expect_true(all(w$per_unit$nobs == 1L))

  l <- suppressWarnings(loo(fit, type = "loso", units = 1:5))
  expect_equal(
    unname(w$estimates["elpd_waic", "Estimate"]),
    l$elpd_2,
    tolerance = 0.05
  )
})