tests/testthat/test-svc-nuts-frontdoor.R

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

# Front-door exact NUTS for spatially-varying coefficients (gcol33/tulpa#158d).
# tulpa(y ~ x, spatial = spatial_svc(~ lon + lat, terms = ~ x - 1), mode =
# "exact") routes through the generic ModelData sampler, which allocates the
# per-term (sigma2, phi) hypers and the [n_svc x n_obs] NNGP field (keyed on
# data.has_svc) and samples them jointly. eta_i += sum_j X_svc[i,j] w_j(s_i),
# so the fixed slope absorbs the mean of the varying coefficient and the SVC
# field carries the spatial variation. SVC is exact-only (no nested front door).

make_svc_pois <- function(n = 120L, seed = 404L) {
  set.seed(seed)
  lon <- runif(n); lat <- runif(n)
  bsurf <- 0.9 * sin(2.8 * lon) + 0.7 * cos(2.2 * lat)   # varying part of x-slope
  x <- rnorm(n)
  d <- data.frame(lon = lon, lat = lat, x = x)
  d$y <- rpois(n, exp(0.2 + (0.8 + bsurf) * x))
  list(d = d, bsurf = bsurf)
}

test_that("svc requires an exact mode (nested/laplace refuse)", {
  s <- make_svc_pois(n = 30L)
  expect_error(
    tulpa(y ~ x, data = s$d, family = "poisson",
          spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L),
          mode = "structured"),
    "Spatially-varying"
  )
})

# HSGP-approximated SVC (gcol33/tulpa#813): the same front door, a Hilbert-
# space basis field instead of an NNGP one. compute_param_layout keys the field
# size on data.svc_hsgp_data.m_total (m^2, shared across terms) rather than
# n_obs, so the stored svc_w[] columns are basis coefficients (beta), not
# per-observation field values -- reconstructing the surface needs the same
# cpp_hsgp_field_predict() krige the HSGP spatial-field predict path uses
# (R/methods_generic.R:.hsgp_field_at), treating each posterior draw as its own
# equally-weighted grid cell.
.hsgp_svc_field_from_draws <- function(fit, spatial, term = 1L) {
  m    <- as.integer(spatial$m)
  Mtot <- m * m
  wcol <- grep("^svc_w\\[", colnames(fit$draws))
  beta_grid <- fit$draws[, wcol[seq((term - 1L) * Mtot + 1L, term * Mtot)],
                        drop = FALSE]
  sigma2_grid     <- exp(fit$draws[, sprintf("log_sigma2_svc[%d]", term)])
  lengthscale_grid <- exp(fit$draws[, sprintf("log_phi_svc[%d]", term)])
  n_draws <- nrow(fit$draws)
  cm <- as.matrix(spatial$coords_matrix)
  as.numeric(cpp_hsgp_field_predict(
    cm, cm, m, as.numeric(spatial$c_boundary),
    beta_grid, sigma2_grid, lengthscale_grid, rep(1 / n_draws, n_draws)))
}

# Rows sharing a site: the NNGP field lives on the distinct locations and each
# row reads its own. Refused before -- as rows, two observations at one site are
# a distance-0 pair, perfectly correlated, and the neighbour covariance singular.
make_svc_sites <- function(n_site = 30L, per = 3L, seed = 405L) {
  set.seed(seed)
  lon <- runif(n_site); lat <- runif(n_site)
  bsurf <- 0.9 * sin(2.8 * lon) + 0.7 * cos(2.2 * lat)
  site <- rep(seq_len(n_site), each = per)
  d <- data.frame(lon = lon[site], lat = lat[site], x = rnorm(n_site * per))
  d$y <- rpois(nrow(d), exp(0.2 + (0.8 + bsurf[site]) * d$x))
  list(d = d, bsurf = bsurf, site = site)
}

skip_on_cran_budget()

test_that("auto mode routes SVC / TVC to the exact ModelData backend (#222)", {
  skip_on_cran_budget()
  sel_svc <- tulpa:::auto_select_mode(
    family = list(name = "poisson"), n_obs = 100L,
    has_spatial = TRUE, has_temporal = FALSE, has_latent = FALSE,
    spatial_type = "svc")
  expect_equal(sel_svc$mode, "exact")
  expect_equal(tulpa:::BACKEND_REGISTRY[[sel_svc$backend]]$input, "modeldata")

  sel_tvc <- tulpa:::auto_select_mode(
    family = list(name = "poisson"), n_obs = 100L,
    has_spatial = FALSE, has_temporal = TRUE, has_latent = FALSE,
    temporal = list(type = "tvc"))
  expect_equal(sel_tvc$mode, "exact")
  expect_equal(tulpa:::BACKEND_REGISTRY[[sel_tvc$backend]]$input, "modeldata")
})

test_that("svc(approx = 'nngp') puts the field on the distinct locations", {
  skip_on_cran_budget()
  s <- make_svc_sites()
  X <- model.matrix(~ x, s$d)
  v <- validate_svc(spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L), s$d, X)
  expect_identical(v$n_obs, nrow(s$d))
  expect_identical(v$n_loc, 30L)
  expect_identical(v$obs_to_loc, s$site)
  expect_identical(nrow(v$neighbor_info$nn_idx), 30L)
  expect_identical(v$n_spatial, 30L)
  spec <- tulpa:::.svc_sampler_spec(v, X)
  expect_identical(nrow(spec$coords), 30L)
  expect_identical(spec$obs_to_loc, s$site - 1L)
  expect_length(spec$X_svc, nrow(s$d))            # still one row per observation
  # Distinct coordinates keep the identity map.
  s1 <- make_svc_pois(n = 30L)
  v1 <- validate_svc(spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L),
                     s1$d, model.matrix(~ x, s1$d))
  expect_null(v1$obs_to_loc)
  expect_null(tulpa:::.svc_sampler_spec(v1, model.matrix(~ x, s1$d))$obs_to_loc)
})

test_that("rows sharing a site read one field value in eta", {
  skip_on_cran_budget()
  s <- make_svc_sites(n_site = 12L, per = 2L)
  X <- model.matrix(~ x, s$d)
  v <- validate_svc(spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 4L,
                                parameterization = "centered"), s$d, X)
  spec <- tulpa:::.svc_sampler_spec(v, X)
  set.seed(9)
  w <- rnorm(12L)
  draw <- matrix(c(0.2, 0.8, log(0.5), log(0.3), w), nrow = 1L)
  eta <- tulpa:::cpp_tulpa_glmm_eta_draws(
    draw, s$d$y, rep(1L, nrow(s$d)), X, "poisson", svc_spec = spec)
  wc <- w - mean(w)                                 # the level is centred away
  expect_equal(as.numeric(eta),
               as.numeric(X %*% c(0.2, 0.8) + s$d$x * wc[s$site]),
               tolerance = 1e-12)
})

skip_on_cran()

test_that("svc fits under the default (auto) mode without erroring (#222)", {
  skip_if_not_slow()
  s <- make_svc_pois(n = 40L)
  # No mode = argument: the default "auto" must route to exact NUTS, not fall to
  # a Laplace/size heuristic that the varying-coefficient guard rejects.
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L),
               control = list(n_iter = 120L, n_warmup = 60L, seed = 1L))
  expect_true(any(grepl("^svc_w\\[", colnames(fit$draws))))
})

test_that("svc exact NUTS allocates the per-term field (structural)", {
  skip_if_not_slow()
  s <- make_svc_pois(n = 40L)
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L),
               mode = "exact",
               control = list(n_iter = 120L, n_warmup = 60L, seed = 1L))
  pn <- colnames(fit$draws)
  expect_true(all(c("log_sigma2_svc[1]", "log_phi_svc[1]") %in% pn))
  expect_equal(sum(grepl("^svc_w\\[", pn)), 40L)   # one field per obs, one term
  expect_equal(ncol(fit$draws), 4L + 40L)           # 2 fixed + 2 hypers + field
})

test_that("both svc parameterizations report an exactly centred field", {
  skip_if_not_slow()
  s <- make_svc_pois(n = 40L)
  for (par in c("noncentered", "centered")) {
    fit <- tulpa(y ~ x, data = s$d, family = "poisson",
                 spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L,
                                       parameterization = par),
                 mode = "exact",
                 control = list(n_iter = 120L, n_warmup = 60L, seed = 1L))
    W <- fit$draws[, grep("^svc_w\\[", colnames(fit$draws)), drop = FALSE]
    # The term's level is identified by CENTRING the field on its way into eta,
    # under either parameterization, so every draw's field sums to zero exactly.
    # A penalty on the sum leaves it merely near zero instead -- at
    # s2z_precision(n) the pinned mean carries sd = 1e-3, which this bound
    # separates from the centred field by ten orders of magnitude.
    expect_lt(max(abs(rowMeans(W))), 1e-12)
  }
})

test_that("svc exact NUTS recovers the varying-coefficient surface", {
  skip_if_not_slow()
  s <- make_svc_pois()
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 8L),
               mode = "exact",
               control = list(n_iter = 350L, n_warmup = 175L, seed = 5L))
  wcol <- grep("^svc_w\\[", colnames(fit$draws))
  fm <- colMeans(fit$draws[, wcol, drop = FALSE])   # w_1(s_i): centred varying slope
  # The SVC field recovers the spatial surface; the fixed x-coefficient absorbs
  # its (non-zero) mean, so recovery is on the centred surface.
  expect_gt(cor(fm - mean(fm), s$bsurf - mean(s$bsurf)), 0.8)
})

test_that("svc exact NUTS is divergence-free and recovers the field scale (#144/#215)", {
  skip_if_not_slow()
  s <- make_svc_pois()
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 8L),
               mode = "exact",
               control = list(n_iter = 350L, n_warmup = 175L, seed = 5L))
  # The #144 regression surfaced as a divergence storm on this exact path, not a
  # field-correlation drop -- so guard the divergences directly, the signal that
  # test survived on.
  expect_lte(n_divergent(fit), ceiling(0.05 * nrow(fit$draws)))   # < 5% divergent
  # Scale recovery: the recovered SVC field SD tracks the true varying-slope SD,
  # and the marginal SD hyperparameter is neither collapsed nor exploded.
  wcol <- grep("^svc_w\\[", colnames(fit$draws))
  sd_ratio <- sd(colMeans(fit$draws[, wcol, drop = FALSE])) / sd(s$bsurf)
  expect_gt(sd_ratio, 0.4)
  expect_lt(sd_ratio, 2.5)
  sigma_svc <- mean(exp(0.5 * fit$draws[, "log_sigma2_svc[1]"]))
  expect_true(is.finite(sigma_svc) && sigma_svc > 0.1)
})

test_that("hsgp svc requires an exact mode (nested/laplace refuse)", {
  skip_on_cran()
  s <- make_svc_pois(n = 30L)
  expect_error(
    tulpa(y ~ x, data = s$d, family = "poisson",
          spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, approx = "hsgp"),
          mode = "structured"),
    "Spatially-varying"
  )
})

test_that("hsgp svc fits under the default (auto) mode without erroring", {
  skip_if_not_slow()
  s <- make_svc_pois(n = 40L)
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1,
                                     approx = "hsgp", m = 5L),
               control = list(n_iter = 120L, n_warmup = 60L, seed = 1L))
  expect_true(any(grepl("^svc_w\\[", colnames(fit$draws))))
})

test_that("hsgp svc exact NUTS allocates m^2 basis coefficients, not n_obs (structural)", {
  skip_if_not_slow()
  s <- make_svc_pois(n = 40L)
  m <- 5L
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1,
                                     approx = "hsgp", m = m),
               mode = "exact",
               control = list(n_iter = 120L, n_warmup = 60L, seed = 1L))
  pn <- colnames(fit$draws)
  expect_true(all(c("log_sigma2_svc[1]", "log_phi_svc[1]") %in% pn))
  expect_equal(sum(grepl("^svc_w\\[", pn)), m^2)   # m^2 basis coefficients, one term
  expect_equal(ncol(fit$draws), 4L + m^2)          # 2 fixed + 2 hypers + basis
})

test_that("hsgp svc exact NUTS recovers the varying-coefficient surface", {
  skip_if_not_slow()
  s <- make_svc_pois()
  spatial <- spatial_svc(~ lon + lat, terms = ~ x - 1, approx = "hsgp", m = 8L)
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial, mode = "exact",
               control = list(n_iter = 350L, n_warmup = 175L, seed = 5L))
  expect_lte(n_divergent(fit), ceiling(0.05 * nrow(fit$draws)))
  spatial <- validate_svc(spatial, s$d, model.matrix(~x, s$d))
  fm <- .hsgp_svc_field_from_draws(fit, spatial)
  expect_gt(cor(fm - mean(fm), s$bsurf - mean(s$bsurf)), 0.8)
})

test_that("svc exact NUTS fits a field with repeated sites", {
  skip_on_cran()
  s <- make_svc_sites()
  fit <- tulpa(y ~ x, data = s$d, family = "poisson",
               spatial = spatial_svc(~ lon + lat, terms = ~ x - 1, nn = 5L),
               mode = "exact",
               control = list(n_iter = 600L, n_warmup = 300L, n_chains = 2L,
                              seed = 1L))
  expect_equal(sum(grepl("^svc_w\\[", colnames(fit$draws))), 30L)
  post <- svc(fit)
  expect_identical(dim(post$draws)[2], nrow(s$d))  # reported per row
  # Rows at one site carry the same draws.
  expect_identical(post$draws[, 1L, 1L], post$draws[, 2L, 1L])
  sm <- summary(post)
  site_mean <- tapply(sm$mean, s$site, `[`, 1L)
  expect_gt(cor(site_mean, s$bsurf - mean(s$bsurf)), 0.5)
})