tests/testthat/test-progress-threads.R

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

# Progress-line outer-thread count (gcol33/tulpa#88).
#
# The shared fit progress reporter appends "| N threads" to its console line
# when the outer grid runs N cells at once, so "ran on N cores" is a property
# of the log itself. Two reporters mirror the same wire format:
#   * the R-side `tulpa_iter_progress()` (counted R loops, e.g. EM-Laplace),
#     which takes the count via its `threads` argument;
#   * the C++ `tulpa_progress::GridProgress` (the outer nested-Laplace grid and
#     the parallel NUTS sampler), which derives it from the realised outer
#     width stamped by `set_width()` inside `run_nested_laplace_grid`.
# A serial loop leaves the count at 1 and the field is omitted -- non-threaded
# callers are byte-for-byte unchanged.

# --------------------------------------------------------------------------- #
# 1. R-side reporter: optional threads field                                  #
# --------------------------------------------------------------------------- #

test_that("tulpa_iter_progress appends the threads field when threads > 1", {
  old <- options(tulpa.nl_progress = list(progress = TRUE,
                                          progress_every = 1L,
                                          progress_throttle = 0))
  on.exit(options(old), add = TRUE)
  out <- capture.output({
    prg <- tulpa_iter_progress("test", total = 3L,
                                        unit = "cells", threads = 8L)
    prg$tick(); prg$tick(); prg$tick(); prg$finish()
  })
  expect_true(any(grepl("\\| 8 threads", out, fixed = FALSE)))
  # The rest of the line is the unchanged format.
  expect_true(any(grepl("^\\[test\\] 3/3 cells \\(100%\\) .* \\| 8 threads$",
                        out)))
})

test_that("tulpa_iter_progress omits the threads field when serial", {
  old <- options(tulpa.nl_progress = list(progress = TRUE,
                                          progress_every = 1L,
                                          progress_throttle = 0))
  on.exit(options(old), add = TRUE)
  out <- capture.output({
    prg <- tulpa_iter_progress("test", total = 3L, unit = "cells")
    prg$tick(); prg$tick(); prg$tick(); prg$finish()
  })
  expect_false(any(grepl("threads", out)))
  expect_true(any(grepl("^\\[test\\] 3/3 cells \\(100%\\)", out)))
})

# --------------------------------------------------------------------------- #
# 2. C++ outer grid: nested-Laplace-joint console line                        #
# --------------------------------------------------------------------------- #

.chain_adj_88 <- function(n_s) {
  nbr <- lapply(seq_len(n_s),
                function(s) setdiff(c(s - 1L, s + 1L), c(0L, n_s + 1L)))
  n_neighbors <- vapply(nbr, length, integer(1))
  list(adj_row_ptr = as.integer(c(0L, cumsum(n_neighbors))),
       adj_col_idx = as.integer(unlist(nbr)) - 1L,
       n_neighbors = as.integer(n_neighbors),
       n_spatial_units = n_s)
}

.fit_joint_88 <- function(n_threads_outer, seed = 11L) {
  set.seed(seed)
  N <- 160L; n_s <- 16L
  adj <- .chain_adj_88(n_s)
  spatial_idx <- sample.int(n_s, N, replace = TRUE)
  rw  <- cumsum(rnorm(n_s, 0, 1 / sqrt(n_s)))
  phi <- rw - mean(rw)
  x   <- rnorm(N)
  Xocc <- cbind(1, x)
  eta_occ <- as.numeric(Xocc %*% c(-0.2, 0.4)) + phi[spatial_idx]
  occur   <- rbinom(N, 1, plogis(eta_occ))
  is_pos  <- occur == 1L
  Xpos    <- Xocc[is_pos, , drop = FALSE]
  spi_pos <- spatial_idx[is_pos]
  eta_pos <- as.numeric(Xpos %*% c(0.1, -0.3)) + 1.5 * phi[spi_pos]
  y_pos   <- rnorm(sum(is_pos), eta_pos, 0.3)

  arm_occ <- list(y = as.numeric(occur), n_trials = rep(1L, N),
                  X = Xocc, spatial_idx = spatial_idx,
                  re_idx = rep(0, N), n_re_groups = 0L, sigma_re = 1.0,
                  family = "binomial", phi = 1.0)
  arm_pos <- list(y = y_pos, n_trials = rep(1L, length(y_pos)),
                  X = Xpos, spatial_idx = spi_pos,
                  re_idx = rep(0, length(y_pos)), n_re_groups = 0L,
                  sigma_re = 1.0, family = "gaussian", phi = 0.09)
  prior <- list(type = "icar", n_spatial_units = adj$n_spatial_units,
                adj_row_ptr = adj$adj_row_ptr, adj_col_idx = adj$adj_col_idx,
                n_neighbors = adj$n_neighbors, sigma_grid = c(0.6, 1.0, 1.5))

  # Both halves are returned: the console line is what this file asserts on,
  # and the fit's own `n_threads_outer_realised` is what the line has to agree
  # with. Reading the width off the fit rather than restating the request keeps
  # the assertion an identity under any clamp (gcol33/tulpa#855).
  out <- capture.output(
    fit <- tulpa_nested_laplace_joint(
      responses = list(
        occ = arm_occ,
        pos = modifyList(arm_pos, list(
          field_coef = list(name = "alpha", grid = c(0.5, 1.0, 1.5))))
      ),
      prior = prior,
      control = list(force_sparse = TRUE,
                     n_threads_outer = n_threads_outer,
                     var_of_means_consistency = FALSE,
                     progress = TRUE, progress.throttle = 0)
    )
  )
  list(out = out, fit = fit)
}

skip_on_cran()

test_that("nested-laplace-joint progress line shows the outer-thread count", {
  skip_on_cran()
  # The count is the REALISED outer width, which is what this file's own header
  # says the C++ reporter derives it from -- and the width the engine is free to
  # resolve below the request: the memory clamp lowers it, and
  # `_R_CHECK_LIMIT_CORES_` caps the team at two, which `src/omp_threads.h` and
  # `cran-comments.md` both state as intended. Measured on this fixture at a
  # request of 4: `| 4 threads` with the variable unset, `| 2 threads` with it
  # set. Asserting the REQUEST made the block fail under `R CMD check` on a tree
  # whose fix was elsewhere (gcol33/tulpa#855).
  #
  # The identity is available rather than a range: the same width the reporter
  # is stamped with travels out of the driver as
  # `out["n_threads_outer_realised"]` and onto the fit
  # (`.nl_attach_outer_threads()`), so the line and the fit can be held to each
  # other and neither can drift from the clamp.
  n_out <- 4L
  r <- .fit_joint_88(n_threads_outer = n_out)
  skip_if(identical(r$fit$n_threads_outer_realised, 1L),
          "the outer width resolved to 1 (OMP_NUM_THREADS or the OpenMP thread limit)")
  joint_lines <- grep("^\\[nested-laplace-joint\\]", r$out, value = TRUE)
  expect_gt(length(joint_lines), 0L)
  thr <- regmatches(joint_lines, regexpr("\\| [0-9]+ threads$", joint_lines))
  expect_gt(length(thr), 0L)
  n_shown <- unique(as.integer(sub("\\| ([0-9]+) threads$", "\\1", thr)))
  expect_length(n_shown, 1L)
  expect_identical(n_shown, r$fit$n_threads_outer_realised)
  # And the regime, so the identity is not satisfied by a serial run: the
  # suffix is printed only above one thread, and the request is the ceiling.
  expect_gt(n_shown, 1L)
  expect_lte(n_shown, n_out)
})

test_that("serial nested-laplace-joint fit omits the threads field", {
  skip_on_cran()
  r <- .fit_joint_88(n_threads_outer = 1L)
  joint_lines <- grep("^\\[nested-laplace-joint\\]", r$out, value = TRUE)
  expect_gt(length(joint_lines), 0L)
  expect_false(any(grepl("threads", joint_lines)))
  # The omission is what a width of one MEANS, not a separate decision: the
  # suffix is printed only above one thread, so the fit has to report one too
  # or the line and the fit are saying different things.
  expect_identical(r$fit$n_threads_outer_realised, 1L)
})