tests/testthat/test-checkpoint-universal.R

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

# Universal grid-cell / per-unit checkpoint/resume (gcol33/tulpa#50).
#
# The joint nested-Laplace path has its own file
# (test-nested-laplace-joint-checkpoint.R). This covers the other fitters the
# checkpoint was extended to: the single-block nested-Laplace grid kernels (all
# of which share run_multi_block_nested_laplace, so ICAR exercises the path),
# the RE-covariance CCD integration, and the per-chain NUTS producer. For each:
# a checkpointed fit equals an un-checkpointed one, a resume loads completed
# units and re-appends nothing, a torn tail from a killed write is discarded and
# re-solved to the same answer, and a file written for different inputs is
# rejected (fingerprint mismatch).

skip_on_cran()

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

# --- single-block nested Laplace (ICAR) -------------------------------------

test_that("single-block nested-Laplace checkpoint == un-checkpointed + resumes", {
  set.seed(11L)
  n_s <- 20L; N <- 160L
  adj <- .ck_chain_adj(n_s)
  w <- as.numeric(arima.sim(n = n_s, list(ar = 0.6))) * 0.7; w <- w - mean(w)
  sidx <- sample(n_s, N, replace = TRUE); x <- rnorm(N)
  y <- rbinom(N, 1, plogis(0.3 + 0.5 * x + w[sidx]))
  X <- cbind(1, x)
  prior <- c(list(type = "icar", tau_grid = 1 / c(0.5, 0.8, 1.2, 1.8)^2,
                  spatial_idx = sidx), adj)
  fit <- function(ck = NULL) {
    ctrl <- list(diagnose_k = FALSE)
    if (!is.null(ck)) ctrl$checkpoint <- ck
    tulpa_nested_laplace(y = y, n_trials = rep(1L, N), X = X, prior = prior,
                         family = "binomial", control = ctrl)
  }
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)

  f_plain <- fit()
  f_ck    <- fit(list(path = path, resume = FALSE))
  expect_equal(f_ck$log_marginal, f_plain$log_marginal, tolerance = 1e-9)
  expect_true(file.exists(path))
  sz <- file.size(path)

  f_re <- fit(list(path = path, resume = TRUE))
  expect_equal(f_re$log_marginal, f_plain$log_marginal, tolerance = 1e-9)
  expect_equal(file.size(path), sz)  # resume re-appended nothing

  b <- readBin(path, "raw", n = file.size(path))
  writeBin(b[seq_len(as.integer(length(b) * 0.6))], path)
  f_torn <- fit(list(path = path, resume = TRUE))
  expect_equal(f_torn$log_marginal, f_plain$log_marginal, tolerance = 1e-9)

  y2 <- rbinom(N, 1, plogis(0.3 + 0.5 * x + w[sidx] + 0.6))
  expect_error(
    tulpa_nested_laplace(y = y2, n_trials = rep(1L, N), X = X, prior = prior,
                         family = "binomial",
                         control = list(diagnose_k = FALSE,
                                        checkpoint = list(path = path, resume = TRUE))),
    "fingerprint")
})

# The equivalence above reads the fields `LaplaceResult` carries, which is
# exactly why it could not see gcol33/tulpa#646: the per-row predictive variance
# lives in the grid driver's own buffer, so `ckpt_serialize` has no way to carry
# it and a loaded cell used to come back holding the buffer's initial 0. A 0
# there is not an absent value -- a caller marginalising the
# (fitted_eta, fitted_eta_var) pair as a grid mixture reads it as a cell of zero
# predictive variance and understates the reported variance in proportion to the
# resumed mass. Both the all-loaded and the partly-loaded grid are checked: the
# torn tail is what a killed run actually comes back as, and it is the case
# where loaded and freshly solved cells have to agree with each other.
test_that("a resumed cell reports the predictive variance, not the buffer's 0", {
  set.seed(11L)
  n_s <- 20L; N <- 160L
  adj <- .ck_chain_adj(n_s)
  w <- as.numeric(arima.sim(n = n_s, list(ar = 0.6))) * 0.7; w <- w - mean(w)
  sidx <- sample(n_s, N, replace = TRUE); x <- rnorm(N)
  y <- rbinom(N, 1, plogis(0.3 + 0.5 * x + w[sidx]))
  X <- cbind(1, x)
  prior <- c(list(type = "icar", tau_grid = 1 / c(0.5, 0.8, 1.2, 1.8)^2,
                  spatial_idx = sidx), adj)
  fit <- function(ck = NULL) {
    ctrl <- list(diagnose_k = FALSE)
    if (!is.null(ck)) ctrl$checkpoint <- ck
    tulpa_nested_laplace(y = y, n_trials = rep(1L, N), X = X, prior = prior,
                         family = "binomial", control = ctrl)
  }
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)

  f_plain <- fit()
  v0 <- f_plain$fitted_eta_var
  expect_false(is.null(v0))
  # The uninterrupted fit is the reference, so it has to carry a variance in
  # every cell for the comparison below to be able to fail.
  expect_true(all(is.finite(v0)))
  expect_false(any(apply(v0, 1L, function(r) all(r == 0))))

  invisible(fit(list(path = path, resume = FALSE)))
  f_all <- fit(list(path = path, resume = TRUE))   # every cell loaded
  expect_false(any(apply(f_all$fitted_eta_var, 1L, function(r) all(r == 0))))
  expect_equal(f_all$fitted_eta_var, v0, tolerance = 1e-12)

  b <- readBin(path, "raw", n = file.size(path))
  writeBin(b[seq_len(as.integer(length(b) * 0.6))], path)
  f_mix <- fit(list(path = path, resume = TRUE))   # loaded + freshly solved
  expect_false(any(apply(f_mix$fitted_eta_var, 1L, function(r) all(r == 0))))
  expect_equal(f_mix$fitted_eta_var, v0, tolerance = 1e-12)
})

# --- RE-covariance nested CCD -----------------------------------------------

test_that("RE-cov nested CCD checkpoint == un-checkpointed + resumes", {
  set.seed(1L)
  G <- 40L; npg <- 10L; N <- G * npg
  grp <- rep(seq_len(G), each = npg); x <- rnorm(N)
  X <- cbind(1, x); Z <- cbind(1, x)
  Sigma <- matrix(c(0.64, 0.24, 0.24, 0.36), 2)
  u <- t(t(chol(Sigma)) %*% matrix(rnorm(2 * G), 2))
  eta <- as.numeric(X %*% c(-0.3, 0.7)) + rowSums(Z * u[grp, ])
  y <- rbinom(N, 1L, plogis(eta))
  re_term <- list(idx = grp, n_groups = G, n_coefs = 2L, Z = Z)
  fit <- function(ck = NULL)
    tulpa_re_cov_nested(y, rep(1L, N), X, re_term, family = "binomial",
                        control = list(diagnose_k = FALSE, checkpoint = ck))
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)

  f0 <- fit()
  f1 <- fit(list(path = path, resume = FALSE))
  expect_equal(f1$weights, f0$weights, tolerance = 1e-9)
  expect_equal(f1$posterior$mean, f0$posterior$mean, tolerance = 1e-9)
  f2 <- fit(list(path = path, resume = TRUE))
  expect_equal(f2$posterior$mean, f0$posterior$mean, tolerance = 1e-9)

  y2 <- rbinom(N, 1L, plogis(eta + 0.6))
  expect_error(
    tulpa_re_cov_nested(y2, rep(1L, N), X, re_term, family = "binomial",
                        control = list(diagnose_k = FALSE,
                                       checkpoint = list(path = path, resume = TRUE))),
    "fingerprint")
})

# --- per-chain NUTS ---------------------------------------------------------

test_that("per-chain NUTS checkpoint reproduces draws and resumes after a crash", {
  set.seed(7L)
  N <- 120L; p <- 3L
  X <- cbind(1, matrix(rnorm(N * (p - 1)), N))
  y <- as.numeric(X %*% c(0.5, -0.8, 0.4) + rnorm(N, sd = 0.7))
  run <- function(ck = "")
    cpp_tulpa_fit_generic_chains(y, X, n_chains = 4L, n_iter = 300L,
                                 n_warmup = 150L, seed = 123L, verbose = FALSE,
                                 checkpoint_path = ck)
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)

  f_plain <- run("")
  f_ck    <- run(path)
  # Determinism: per-chain checkpoint must not perturb the draws.
  expect_equal(f_ck$draws, f_plain$draws, tolerance = 1e-10)

  # Crash mid-run: keep half the records, resume re-runs only the lost chains.
  b <- readBin(path, "raw", n = file.size(path))
  writeBin(b[seq_len(as.integer(length(b) * 0.5))], path)
  f_re <- run(path)
  expect_equal(f_re$draws, f_plain$draws, tolerance = 1e-10)

  # A full resume loads every chain and appends nothing.
  sz <- file.size(path)
  f_re2 <- run(path)
  expect_equal(file.size(path), sz)
  expect_equal(f_re2$draws, f_plain$draws, tolerance = 1e-10)

  # Different seed -> different fit -> fingerprint mismatch.
  expect_error(
    cpp_tulpa_fit_generic_chains(y, X, n_chains = 4L, n_iter = 300L,
                                 n_warmup = 150L, seed = 999L, verbose = FALSE,
                                 checkpoint_path = path),
    "fingerprint")
})

# --- per-chain NUTS, reached from the R door (gcol33/tulpa#808) -------------
#
# The block above exercises the checkpoint mechanism through the unexported
# Gaussian test entry cpp_tulpa_fit_generic_chains(); this block exercises the
# same mechanism through the user-reachable route -- tulpa(mode = "hmc") /
# tulpa_sample_glmm() -- which previously accepted control$checkpoint and
# dropped it silently.

test_that("tulpa(mode = 'hmc') checkpoint reproduces draws and resumes", {
  set.seed(8L)
  n <- 100L
  d <- data.frame(x = rnorm(n))
  d$y <- rpois(n, exp(0.3 + 0.5 * d$x))
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)

  run <- function(ck = NULL) {
    ctrl <- list(n_iter = 200L, warmup = 100L, n_chains = 2L, seed = 1L)
    if (!is.null(ck)) ctrl$checkpoint <- list(path = ck, resume = TRUE)
    tulpa(y ~ x, data = d, family = "poisson", mode = "hmc", control = ctrl)
  }

  f_plain <- run()
  f_ck    <- run(path)
  expect_true(file.exists(path))
  expect_equal(f_ck$draws, f_plain$draws, tolerance = 1e-10)

  # Resume: the file already holds every finished chain, so a second call
  # against the same path reproduces the draws without appending.
  sz <- file.size(path)
  f_re <- run(path)
  expect_equal(file.size(path), sz)
  expect_equal(f_re$draws, f_plain$draws, tolerance = 1e-10)

  # resume = FALSE removes the stale file before the kernel opens it, so a
  # fresh run against different data does not inherit the old chains.
  d2 <- data.frame(x = rnorm(n))
  d2$y <- rpois(n, exp(-0.2 + 0.1 * d2$x))
  f_fresh <- tulpa(y ~ x, data = d2, family = "poisson", mode = "hmc",
                   control = list(n_iter = 200L, warmup = 100L, n_chains = 2L,
                                  seed = 1L,
                                  checkpoint = list(path = path, resume = FALSE)))
  expect_false(isTRUE(all.equal(unname(f_fresh$draws), unname(f_plain$draws))))
})

test_that("control$checkpoint is refused on a non-NUTS sampler backend", {
  set.seed(9L)
  n <- 60L
  d <- data.frame(x = rnorm(n))
  d$y <- rpois(n, exp(0.2 + 0.3 * d$x))
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)
  expect_error(
    tulpa(y ~ x, data = d, family = "poisson", mode = "vi",
         control = list(checkpoint = list(path = path))),
    "only read by the NUTS/HMC kernel")
  expect_false(file.exists(path))
})

# --- gcol33/tulpa#442 -------------------------------------------------------
#
# A resumed cell has to be indistinguishable from a freshly solved one in EVERY
# field, and a fingerprint has to separate every input that moves a cell.

test_that("a resumed grid reproduces the diagnostics, not only the marginal", {
  set.seed(19L)
  n_s <- 14L; N <- 110L
  adj <- .ck_chain_adj(n_s)
  w <- as.numeric(arima.sim(n = n_s, list(ar = 0.5))) * 0.7; w <- w - mean(w)
  sidx <- sample(n_s, N, replace = TRUE); x <- rnorm(N)
  y <- rbinom(N, 1, plogis(0.2 + 0.6 * x + w[sidx]))
  X <- cbind(1, x)
  prior <- c(list(type = "icar", tau_grid = 1 / c(0.6, 1.0, 1.6)^2,
                  spatial_idx = sidx), adj)
  fit <- function(ck = NULL) {
    ctrl <- list(diagnose_k = FALSE, diagnose_skew = TRUE)
    if (!is.null(ck)) ctrl$checkpoint <- ck
    tulpa_nested_laplace(y = y, n_trials = rep(1L, N), X = X, prior = prior,
                         family = "binomial", control = ctrl)
  }
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)

  f_plain <- fit()
  f_ck    <- fit(list(path = path, resume = FALSE))
  f_re    <- fit(list(path = path, resume = TRUE))

  # The skew payload is what the eleven-field record used to drop: a resumed
  # cell came back with an empty inner_skew and read as "did not compute it".
  expect_equal(f_re$inner_skew, f_plain$inner_skew, tolerance = 1e-12)
  expect_equal(f_re$inner_skew, f_ck$inner_skew, tolerance = 1e-12)
  expect_equal(f_re$log_marginal, f_plain$log_marginal, tolerance = 1e-9)
  expect_equal(f_re$modes, f_plain$modes, tolerance = 1e-9)
})

test_that("the RE group assignment separates two fingerprints in its second half", {
  # re_idx is a NumericVector folded at sizeof(int): only its leading half
  # reached the hash, so two assignments agreeing there matched and a resume
  # replayed one run's cells under the other's data.
  set.seed(23L)
  n_s <- 10L; N <- 80L
  adj <- .ck_chain_adj(n_s)
  sidx <- sample(n_s, N, replace = TRUE); x <- rnorm(N)
  y <- rbinom(N, 1, plogis(0.2 + 0.5 * x))
  X <- cbind(1, x)
  g1 <- rep(c(1L, 2L), each = N / 2)
  g2 <- rep(c(1L, 3L), each = N / 2)   # differs only in the SECOND half
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)
  run <- function(g, resume) {
    prior <- c(list(type = "icar", tau_grid = 1 / c(0.8, 1.4)^2,
                    spatial_idx = sidx), adj)
    tulpa_nested_laplace(y = y, n_trials = rep(1L, N), X = X, prior = prior,
                         family = "binomial", re_idx = g,
                         n_re_groups = max(g), sigma_re = 1,
                         control = list(diagnose_k = FALSE,
                                        checkpoint = list(path = path,
                                                          resume = resume)))
  }
  run(g1, FALSE)
  expect_error(run(g2, TRUE), "fingerprint")
})

# ============================================================================
# The chain fingerprint covers the DATA, not only its shape (gcol33/tulpa#683)
#
# The response lives behind an opaque pointer, so the fingerprint folds the
# target the chains sample -- evaluated at engine-fixed probe positions -- and a
# resume against a different data set of the same shape is refused instead of
# returning the earlier fit's draws under the new data's name.
# ============================================================================

test_that("a resume against different data of the same shape is refused", {
  set.seed(11L)
  X  <- cbind(1, rnorm(60))
  y1 <- rpois(60, exp( 0.5 + 0.4 * X[, 2]))
  y2 <- rpois(60, exp(-1.5 - 1.2 * X[, 2]))
  cp <- file.path(tempdir(), "ckpt-683.bin")
  on.exit(unlink(cp), add = TRUE)
  unlink(cp)
  run <- function(y, path) cpp_tulpa_fit_generic_chains(
    as.numeric(y), X, 2L, n_iter = 120L, n_warmup = 60L, seed = 5L,
    checkpoint_path = path)

  f1 <- run(y1, cp)
  expect_error(run(y2, cp), "fingerprint mismatch")

  # The same data still resumes, bit for bit: the probe is deterministic and
  # takes nothing from R's RNG stream.
  f1b <- run(y1, cp)
  expect_identical(f1$draws, f1b$draws)

  # And it is the DATA that separates them, not the file: y2 fits normally on
  # its own path, to its own answer.
  f2 <- run(y2, file.path(tempdir(), "ckpt-683b.bin"))
  on.exit(unlink(file.path(tempdir(), "ckpt-683b.bin")), add = TRUE)
  expect_false(isTRUE(all.equal(colMeans(f1$draws), colMeans(f2$draws))))
})

test_that("a changed design of the same shape is refused too", {
  set.seed(12L)
  X1 <- cbind(1, rnorm(50))
  X2 <- cbind(1, rnorm(50))
  y  <- rpois(50, 3)
  cp <- file.path(tempdir(), "ckpt-683c.bin")
  on.exit(unlink(cp), add = TRUE)
  unlink(cp)
  run <- function(X, path) cpp_tulpa_fit_generic_chains(
    as.numeric(y), X, 2L, n_iter = 120L, n_warmup = 60L, seed = 3L,
    checkpoint_path = path)
  run(X1, cp)
  expect_error(run(X2, cp), "fingerprint mismatch")
})

# A front-door nested fit whose default outer axis is recentred solves a
# SECOND grid onto the same checkpoint path. The single-block fingerprint used
# to fold the grid's axis values, so that refit refused the file the first grid
# had just written and the fit could not checkpoint at all -- hidden while the
# checkpoint vignette ran mode = "laplace", which dropped `control$checkpoint`
# (gcol33/tulpa#870). Cells are keyed by their exact coordinate, so the two
# grids share one file.
test_that("a recentred front-door nested fit checkpoints and resumes", {
  skip_on_cran()
  set.seed(20260716)
  S <- 30L
  W <- matrix(0, S, S)
  for (i in 1:(S - 1)) W[i, i + 1] <- W[i + 1, i] <- 1
  df <- data.frame(region = factor(seq_len(S)))
  df$x <- as.integer(df$region) / 10 + rnorm(S, 0, 0.3)
  df$y <- rbinom(S, 20, plogis(-0.4 + 0.5 * df$x))
  path <- tempfile(fileext = ".ckpt"); on.exit(unlink(path), add = TRUE)
  fit <- function(resume) suppressMessages(tulpa(
    y ~ x + spatial(region), data = df, family = "binomial",
    n_trials = rep(20L, S), spatial = spatial_car(W, level = "obs"),
    mode = "structured",
    control = list(checkpoint = list(path = path, resume = resume))))
  f1 <- fit(FALSE)
  expect_true(file.exists(path))
  sz <- file.size(path)
  expect_gt(sz, 16)
  f2 <- fit(TRUE)
  expect_equal(coef(f2), coef(f1), tolerance = 1e-9)
  expect_equal(file.size(path), sz)  # the resume re-solved nothing
})