FEM Matérn joint density with adlaplace

library(adlaplace)

## 
## Attaching package: 'adlaplace'

## The following objects are masked from 'package:stats':
## 
##     binomial, gaussian

library(adlaplaceGrf)
library(Matrix)

This vignette builds a tiny tensor-product B-spline FEM Matérn field, wires it into adlaplace via random_fem_ssq_2 (quadratic form) and random_fem_det_2 (log-determinant), and evaluates the joint log-density, gradient, and Hessian. Optimization is deferred.

FEM ingredients

set.seed(1)
knots_list <- list(
  x = seq(0, 1, length.out = 5),
  y = seq(0, 1, length.out = 5)
)
sites_eval <- do.call(expand.grid, knots_list)
fem <- grf_bspline(sites_eval, knots_list, degree = 2L)
prec <- fem_precision_payload(fem, alpha = 2L)
nr <- nrow(fem$C)
c(n_obs = nrow(sites_eval), n_basis = nr, nnz_Q = length(prec$Q_i))

##   n_obs n_basis   nnz_Q 
##      25      36     306

Hierarchical toy model

Gaussian observations \(y \mid w \sim N(Aw, \sigma^2 I)\) with FEM prior on the spline weights \(w\). The prior splits into a random shard (random_fem_ssq_2) and a parameters shard (random_fem_det_2). Hyperparameters are \((\log\sigma, \log\rho, \log\sigma_w)\) where \(\rho\) is the practical range and \(\sigma_w\) is the field SD.

A <- fem$A
n <- nrow(A)
sigma <- 0.4
nu <- 1
kappa <- 2
range <- sqrt(8 * nu) / kappa
sd_w <- 1
w_true <- rnorm(nr, sd = 0.2)
y <- as.numeric(A %*% w_true + rnorm(n, sd = sigma))

# Parameter layout: gamma (nr), theta = (log_sigma, log_range, log_sd)
theta_obs <- Matrix::sparseMatrix(i = 1L, j = 1L, dims = c(3L, 1L))
theta_fem <- Matrix::sparseMatrix(
  i = c(2L, 3L), j = c(1L, 2L), x = c(1, 1), dims = c(3L, 2L)
)

obs <- adlaplace:::ad_data(
  y = y,
  A = A,
  gamma_map = Matrix::Diagonal(nr),
  theta_map = theta_obs,
  ad_kind = "observations",
  ad_fun = "gaussian_obs"
)
extra <- adlaplace::ad_data(
  obs,
  ad_kind = "parameters",
  ad_fun = "gaussian_extra"
)
rand_ssq <- adlaplace:::ad_data(
  gamma_map = Matrix::Diagonal(nr),
  theta_map = theta_fem,
  ad_kind = "random",
  ad_fun = "random_fem_ssq_2",
  package = "adlaplaceGrf",
  precision = prec
)
rand_det <- adlaplace:::ad_data(
  beta_map = 0L,
  gamma_map = nr,
  theta_map = theta_fem,
  ad_kind = "parameters",
  ad_fun = "random_fem_det_2",
  package = "adlaplaceGrf",
  precision = prec
)

config <- list(
  gamma = rep(0, nr),
  theta = c(log(sigma), log(range), log(sd_w)),
  transform_theta = TRUE
)
ptr <- c(
  ad_fun_ptr(obs, config),
  ad_fun_ptr(extra, config),
  ad_fun_ptr(rand_ssq, config),
  ad_fun_ptr(rand_det, config)
)

Joint density, gradient, and Hessian

w <- rnorm(nr, sd = 0.1)
x <- c(w, log(sigma), log(range), log(sd_w))
dens <- joint_log_dens(ptr, x, negative = FALSE)
g <- grad(ptr, x, inner = FALSE, negative = FALSE)
H <- hessian(ptr, x, inner = FALSE, negative = FALSE)
c(
  log_dens = dens,
  grad_norm = sqrt(sum(as.numeric(g)^2)),
  hess_nrow = nrow(H),
  hess_nnz = length(H@x)
)

##  log_dens grad_norm hess_nrow  hess_nnz 
## -38.45517  45.78595  39.00000 418.00000

Finite-difference checks

Along one coordinate (here \(\log\rho\)), compare the AD gradient to a finite difference of the joint log-density, and AD Hessian entries to finite differences of the AD gradient — the same check as in the adlaplace vignette.

par(mfrow = c(2, 2), mar = c(2.5, 2.5, 1.5, 0.5), mgp = c(1.4, 0.5, 0.5))

Dpar <- length(x) - 1L # log range
Ngrid <- 11L
par_grid <- matrix(x, nrow = Ngrid, ncol = length(x), byrow = TRUE)
Sx <- x[Dpar] + seq(-0.15, 0.15, length.out = Ngrid)
SxD <- Sx[-1] - diff(Sx) / 2
par_grid[, Dpar] <- Sx
x_list <- split(par_grid, row(par_grid))

dens_scan <- vapply(
  x_list,
  joint_log_dens,
  numeric(1),
  ad_fun = ptr,
  negative = FALSE
)

grad_scan <- do.call(
  cbind,
  lapply(x_list, grad, ad_fun = ptr, inner = FALSE, negative = FALSE)
)

plot(Sx, dens_scan, type = "l", xlab = expression(log(rho)), ylab = "log dens")
plot(Sx, grad_scan[Dpar, ], type = "l", xlab = expression(log(rho)), ylab = "AD gradient")
points(SxD, diff(dens_scan) / diff(Sx), pch = 16)
legend("topright", c("AD", "finite diff"), lty = c(1, NA), pch = c(NA, 16), bty = "n")

hes_scan <- array(NA_real_, dim = c(length(x), length(x), Ngrid))
for (i in seq_len(Ngrid)) {
  hes_scan[, , i] <- as.matrix(hessian(
    ptr, x_list[[i]],
    inner = FALSE, negative = FALSE
  ))
}
grad_slope <- apply(grad_scan, 1, diff) / mean(diff(Sx))

# Hessian row for log range vs a weight and vs log sd
for (Dpar2 in c(1L, length(x))) {
  plot(
    Sx, hes_scan[Dpar, Dpar2, ],
    type = "l",
    xlab = expression(log(rho)),
    ylab = paste0("H[", Dpar, ",", Dpar2, "]"),
    ylim = range(c(hes_scan[Dpar, Dpar2, ], grad_slope[, Dpar2]), na.rm = TRUE)
  )
  points(SxD, grad_slope[, Dpar2], pch = 16)
}

Pointwise central differences at the base x summarize agreement across all coordinates:

eps <- 1e-5
g_fd <- vapply(seq_along(x), function(i) {
  xp <- xm <- x
  xp[i] <- xp[i] + eps
  xm[i] <- xm[i] - eps
  (joint_log_dens(ptr, xp, negative = FALSE) -
    joint_log_dens(ptr, xm, negative = FALSE)) / (2 * eps)
}, numeric(1))

# FD Hessian column for log range from AD gradients
j <- Dpar
H_fd_j <- vapply(seq_along(x), function(i) {
  xp <- xm <- x
  xp[i] <- xp[i] + eps
  xm[i] <- xm[i] - eps
  gp <- as.numeric(grad(ptr, xp, inner = FALSE, negative = FALSE))
  gm <- as.numeric(grad(ptr, xm, inner = FALSE, negative = FALSE))
  (gp[j] - gm[j]) / (2 * eps)
}, numeric(1))

c(
  max_abs_grad_diff = max(abs(as.numeric(g) - g_fd)),
  max_abs_hess_col_diff = max(abs(as.numeric(H[, j]) - H_fd_j))
)

##     max_abs_grad_diff max_abs_hess_col_diff 
##          1.567795e-09          2.330021e-09

No outer optimization is performed here; that will appear in a later vignette.