tests/testthat/test-nl-grid-cap.R

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

test_that("the guard refuses past a set ceiling and names the axes", {
  expect_silent(.nl_check_grid_cap(1e6, Inf, "Reduce it."))
  expect_silent(.nl_check_grid_cap(2401, 2401, "Reduce it."))
  expect_error(
    .nl_check_grid_cap(24, 4, "Reduce it.", " It crosses b1 (x) x phi_b 4."),
    "24 cells, past control$max_grid_cells = 4. It crosses b1 (x) x phi_b 4. Reduce it.",
    fixed = TRUE)
})

test_that("the crossing names the dispersion axes crossed on top", {
  g <- matrix(1:6, ncol = 1, dimnames = list(NULL, "tau"))
  expect_identical(.nl_grid_crossing(list(g), list(b = c(0.1, 0.4, 1, 3))),
                   " It crosses b1 (6 rows: tau 6) x phi_b 4.")
})

test_that("the timing warning carries the count, the layout and the remedy", {
  over <- .NL_MULTI_GRID_WARN_SECONDS + 1
  expect_warning(
    .nl_multi_grid_warn(over, 10000, "Reduce per-block or phi grid sizes.",
                        " It crosses b1 (50 rows: sigma 5 x alpha 10) x phi_pos 4."),
    paste0("Multi-block outer grid \\(10000 cells\\) took .* to solve\\. It crosses ",
           "b1 \\(50 rows: sigma 5 x alpha 10\\) x phi_pos 4\\. Reduce per-block ",
           "or phi grid sizes\\."))
})

test_that("a dense grid past the threshold is announced before its solve", {
  n <- .NL_DENSE_GRID_WARN_CELLS + 1L
  expect_warning(
    .nl_dense_grid_warn(n, "Reduce per-block grid sizes.",
                        " It crosses b1 (x) x phi_pos 4."),
    paste0("dense tensor of ", n, " cells, one inner solve each\\. It crosses ",
           "b1 \\(x\\) x phi_pos 4\\. Reduce per-block grid sizes\\."))
  expect_silent(.nl_dense_grid_warn(.NL_DENSE_GRID_WARN_CELLS, "x"))
  expect_silent(.nl_with_internal_batch(.nl_dense_grid_warn(n, "x")))
})

# --------------------------------------------------------------------------- #
# (3) Grids past the threshold warn and are solved                             #
# --------------------------------------------------------------------------- #

.cap_iid_data <- function(seed = 11L, N = 40L, n_a = 5L, n_b = 4L) {
  set.seed(seed)
  ia <- rep_len(seq_len(n_a), N)
  ib <- rep_len(seq_len(n_b), N)
  x  <- rnorm(N)
  eta <- -0.2 + 0.5 * x + rnorm(n_a, 0, 0.3)[ia] + rnorm(n_b, 0, 0.3)[ib]
  list(y = rbinom(N, 1L, plogis(eta)), n = rep(1L, N), X = cbind(1, x),
       ia = as.integer(ia), ib = as.integer(ib), n_a = n_a, n_b = n_b)
}

.cap_iid_prior <- function(d, g_a, g_b) {
  list(
    list(type = "iid", obs_idx = d$ia, n_units = d$n_a, sigma_grid = g_a),
    list(type = "iid", obs_idx = d$ib, n_units = d$n_b, sigma_grid = g_b)
  )
}

.cap_chain_adj <- function(n) {
  nb <- lapply(seq_len(n), function(s) setdiff(c(s - 1L, s + 1L), c(0L, n + 1L)))
  list(adj_row_ptr = as.integer(c(0L, cumsum(lengths(nb)))),
       adj_col_idx = as.integer(unlist(nb) - 1L),
       n_neighbors = as.integer(lengths(nb)))
}

.cap_joint_fixture <- function(g_a, g_b, seed = 19L, n_s = 6L, N = 60L) {
  set.seed(seed)
  adjA <- .cap_chain_adj(n_s); adjB <- .cap_chain_adj(n_s)
  iA <- sample.int(n_s, N, replace = TRUE)
  iB <- sample.int(n_s, N, replace = TRUE)
  fA <- as.numeric(scale(cumsum(rnorm(n_s, 0, 0.4))))
  fB <- as.numeric(scale(cumsum(rnorm(n_s, 0, 0.4))))
  x  <- rnorm(N); X <- cbind(1, x)
  eta <- as.numeric(X %*% c(-0.2, 0.4)) + fA[iA] + fB[iB]
  y1 <- rbinom(N, 1L, plogis(eta))
  y2 <- rnorm(N, eta, 0.5)
  responses <- list(
    a = list(y = as.numeric(y1), n_trials = rep(1L, N), X = X,
             spatial_idx = as.integer(iA), re_idx = rep(0, N),
             n_re_groups = 0L, sigma_re = 1.0, family = "binomial", phi = 1.0),
    b = list(y = y2, n_trials = rep(1L, N), X = X,
             spatial_idx = as.integer(iA), re_idx = rep(0, N),
             n_re_groups = 0L, sigma_re = 1.0, family = "gaussian", phi = 0.25))
  prior <- list(
    list(type = "icar", n_spatial_units = n_s,
         adj_row_ptr = adjA$adj_row_ptr, adj_col_idx = adjA$adj_col_idx,
         n_neighbors = adjA$n_neighbors, tau_grid = g_a,
         spatial_idx = list(as.integer(iA), as.integer(iA))),
    list(type = "icar", n_spatial_units = n_s,
         adj_row_ptr = adjB$adj_row_ptr, adj_col_idx = adjB$adj_col_idx,
         n_neighbors = adjB$n_neighbors, tau_grid = g_b,
         spatial_idx = list(as.integer(iB), as.integer(iB))))
  list(responses = responses, prior = prior)
}

test_that("the joint dispatch counts the phi cells in the pre-solve warning", {
  # 46 x 23 = 1058 latent cells, under the threshold alone; crossed with two
  # dispersion values it is 2116 inner solves.
  f <- .cap_joint_fixture(seq(0.5, 20, length.out = 46L),
                          seq(0.5, 20, length.out = 23L))
  msg <- tryCatch(
    tulpa_nested_laplace_joint(responses = f$responses, prior = f$prior,
                               phi_grid = list(b = c(0.1, 0.4)),
                               control = list(diagnose_k = FALSE,
                                              integration = "grid")),
    warning = conditionMessage)
  expect_match(msg, "dense tensor of 2116 cells", fixed = TRUE)
  expect_match(msg, " x phi_b 2.", fixed = TRUE)
  expect_match(msg, "Reduce per-block or phi grid sizes", fixed = TRUE)
})

skip_on_cran_budget()
# test-nl-grid-cap.R
# A multi-block outer grid is solved at whatever size its axes multiply out to
# unless the caller sets `control$max_grid_cells` (gcol33/tulpa#916). A cell
# count does not measure the run's cost -- that is cells times the cost of one
# inner solve -- so by default a large dense grid is announced with a warning
# before its solve, and the warnings and the opt-in refusal name the axes that
# produced the count (gcol33/tulpa#913), dispersion axes included
# (gcol33/tulpa#915).

# --------------------------------------------------------------------------- #
# (1) The opt-in ceiling                                                       #
# --------------------------------------------------------------------------- #

test_that(".nl_max_grid_cells() has no default ceiling and reads the knob", {
  skip_on_cran_budget()
  expect_identical(.nl_max_grid_cells(), Inf)
  expect_identical(.nl_max_grid_cells(list()), Inf)
  expect_identical(.nl_max_grid_cells(list(max_grid_cells = 4096L)), 4096)
  op <- options(tulpa.nl_max_grid_cells = 3000)
  on.exit(options(op), add = TRUE)
  expect_identical(.nl_max_grid_cells(), 3000)
  expect_identical(.nl_max_grid_cells(list(max_grid_cells = 12)), 12)
})

test_that(".nl_max_grid_cells() rejects a malformed knob", {
  skip_on_cran_budget()
  for (v in list(0, -5, c(10, 20), "many", NA_real_))
    expect_error(.nl_max_grid_cells(list(max_grid_cells = v)), "max_grid_cells")
})

test_that("max_grid_cells is an accepted control knob on every front door", {
  skip_on_cran_budget()
  expect_silent(tulpa_check_control(list(max_grid_cells = 4096),
                                    .CONTROL_KEYS$nested_laplace,
                                    "tulpa_nested_laplace"))
  expect_silent(tulpa_check_control(list(max_grid_cells = 4096),
                                    .CONTROL_KEYS$nested_laplace_joint,
                                    "tulpa_nested_laplace_joint"))
  expect_silent(tulpa_check_control(list(max_grid_cells = 4096),
                                    .CONTROL_KEYS$tulpa, "tulpa"))
})

# --------------------------------------------------------------------------- #
# (2) The layout the timing warning carries                                    #
# --------------------------------------------------------------------------- #

test_that("the layout names each block's rows and per-axis levels", {
  skip_on_cran_budget()
  g1 <- as.matrix(expand.grid(sigma = c(0.1, 1, 3), alpha = c(0, 0.5, 1, 2)))
  g2 <- matrix(c(1, 2), ncol = 1, dimnames = list(NULL, "tau"))
  expect_identical(.nl_grid_layout(list(g1, g2)),
                   "b1 (12 rows: sigma 3 x alpha 4) x b2 (2 rows: tau 2)")
  expect_identical(.nl_grid_crossing(list(g1, g2)),
                   " It crosses b1 (12 rows: sigma 3 x alpha 4) x b2 (2 rows: tau 2).")
  expect_identical(.nl_grid_crossing(), "")
})

# The warning is raised before the first inner solve, so catching it as a
# condition stops the fit there and nothing below is fitted.
test_that("the multi-block dispatch warns before solving a large dense grid", {
  skip_on_cran_budget()
  d <- .cap_iid_data()
  prior <- .cap_iid_prior(d, seq(0.05, 2, length.out = 46L),
                             seq(0.05, 2, length.out = 46L))       # 2116 cells
  msg <- tryCatch(
    tulpa_nested_laplace(y = d$y, n_trials = d$n, X = d$X, prior = prior,
                         family = "binomial"),
    warning = conditionMessage)
  expect_match(msg, "dense tensor of 2116 cells", fixed = TRUE)
  expect_match(msg, "It crosses b1 (46 rows: ", fixed = TRUE)
})

test_that("a set ceiling reaches the multi-block dispatch", {
  skip_on_cran_budget()
  d <- .cap_iid_data()
  prior <- .cap_iid_prior(d, c(0.2, 0.5, 1.0), c(0.2, 0.5, 1.0))   # 9 cells
  msg <- tryCatch(
    tulpa_nested_laplace(y = d$y, n_trials = d$n, X = d$X, prior = prior,
                         family = "binomial", control = list(max_grid_cells = 4)),
    error = conditionMessage)
  expect_match(msg, "9 cells, past control$max_grid_cells = 4.", fixed = TRUE)
  expect_match(msg, "It crosses b1 (3 rows: ", fixed = TRUE)
})

test_that("a set ceiling on the joint dispatch counts the phi cells", {
  skip_on_cran_budget()
  # 2 x 2 latent cells fit under a ceiling of 4; crossed with three dispersion
  # values they are 12 inner solves.
  f <- .cap_joint_fixture(c(0.5, 2.0), c(0.5, 2.0))
  msg <- tryCatch(
    tulpa_nested_laplace_joint(responses = f$responses, prior = f$prior,
                               phi_grid = list(b = c(0.1, 0.25, 0.6)),
                               control = list(diagnose_k = FALSE,
                                              integration = "grid",
                                              max_grid_cells = 4)),
    error = conditionMessage)
  expect_match(msg, "12 cells, past control$max_grid_cells = 4.", fixed = TRUE)
  expect_match(msg, " x phi_b 3.", fixed = TRUE)
  expect_match(msg, "Reduce per-block or phi grid sizes", fixed = TRUE)
})

skip_on_cran()

test_that("the multi-block dispatch integrates a 2116-cell grid end to end", {
  skip_on_cran()
  d <- .cap_iid_data()
  prior <- .cap_iid_prior(d, seq(0.05, 2, length.out = 46L),
                             seq(0.05, 2, length.out = 46L))       # 2116 cells
  fit <- suppressWarnings(tulpa_nested_laplace(
    y = d$y, n_trials = d$n, X = d$X, prior = prior, family = "binomial",
    control = list(max_iter = 30L, tol = 1e-6, progress = FALSE)))
  expect_equal(length(fit$weights), 2116L)
  expect_equal(sum(fit$weights), 1)
  expect_true(all(is.finite(fit$theta_mean)))
})

test_that("the joint dispatch solves every latent x phi cell of the dense tensor", {
  skip_on_cran()
  # 2 x 2 latent cells crossed with three dispersion values: 12 inner solves.
  f <- .cap_joint_fixture(c(0.5, 2.0), c(0.5, 2.0))
  fit <- tulpa_nested_laplace_joint(
    responses = f$responses, prior = f$prior,
    phi_grid = list(b = c(0.1, 0.25, 0.6)),
    control = list(diagnose_k = FALSE, integration = "grid",
                   max_iter = 60L, tol = 1e-6))
  expect_s3_class(fit, "tulpa_nested_laplace_joint_multi")
  # Refinement nodes come on top of the tensor, which is solved whole.
  expect_equal(sum(!nzchar(fit$refining_axis)), 12L)
  expect_gte(length(fit$weights), 12L)
})