--- title: "Laplace Approximations with adlaplace" author: "Patrick Brown" date: "`r Sys.Date()`" output: rmarkdown::html_vignette: toc: true vignette: > %\VignetteIndexEntry{Laplace Approximations with adlaplace} %\VignetteEngine{knitr::rmarkdown} %\VignetteEncoding{UTF-8} --- ```{r setup, include=FALSE} library("Matrix") library("adlaplace") # 2 threads when this build has OpenMP; otherwise serial. num_threads <- if (adlaplace::has_openmp()) 2L else 1L ``` This vignette walks through a hierarchical negative-binomial model: simulate data, build an automatic-differentiation (AD) representation as **shards**, validate derivatives, solve the inner Laplace problem for random effects, and fit outer parameters with a profile likelihood. ## Model Observations follow a negative binomial with mean $\mu = \exp(\eta)$. We use two equivalent parameterizations: NB **size** $r$, and Gamma-mixing **standard deviation** $\theta$ (with $E[Z]=1$, $\mathrm{sd}(Z)=\theta$ in the Poisson–Gamma construction). They are linked by $r = 1/\theta^2$. **Size $r$.** $$ \begin{aligned} f(y;\mu,r) = & \log\Gamma(y + r) - \log\Gamma(r) - \log\Gamma(y + 1) + r \log\left(\frac{r}{r + \mu}\right) + y \log\left(\frac{\mu}{r + \mu}\right) \\ = & \log\Gamma(y + r) - \log\Gamma(r) - \log\Gamma(y + 1) + r \log r - r \log(r + \mu) + y \log(\mu) - y \log(r + \mu) \end{aligned} $$ **SD $\theta$** ($r = 1/\theta^2$). $$ \begin{aligned} f(y;\mu,\theta) = & \log\Gamma\left(y + \theta^{-2}\right) - \log\Gamma\left(\theta^{-2}\right) - \log\Gamma(y + 1) + \theta^{-2} \log\left(\frac{\theta^{-2}}{\theta^{-2} + \mu}\right) + y \log\left(\frac{\mu}{\theta^{-2} + \mu}\right) \\ = & \log\Gamma\left(y + \theta^{-2}\right) - \log\Gamma\left(\theta^{-2}\right) - \log\Gamma(y + 1) - 2 \theta^{-2} \log \theta - \theta^{-2} \log\left(\theta^{-2} + \mu\right) + y \log(\mu) - y \log\left(\theta^{-2} + \mu\right) \end{aligned} $$ The linear predictor is $\eta = X\beta + A\gamma$, with two independent random-effect groups in $A$ and log-standard-deviation parameters for each group plus a global overdispersion parameter. Densities are implemented in `src/densities.cpp` with headers under `src/include/adlaplace/densities.hpp` (e.g. negative binomial observation and hyperparameter terms, plus `random_diagonal` for random effects). ## Simulated data We generate 5000 observations, fixed effects, two random-effect groups, and Poisson–Gamma overdispersion. Outer parameters are stored on the log-SD scale (`transform_theta = TRUE`). ```{r theData} Nobs <- 5000 Nrandom1 <- 10 Nrandom2 <- 25 set.seed(0) X <- Matrix::Matrix(cbind(1, rbinom(Nobs, 1, prob = 0.5))) Adf <- data.frame( r1 = sample(Nrandom1, Nobs, replace = TRUE), r2 = sample(Nrandom2, Nobs, replace = TRUE) ) AmatList <- list( r1 = Matrix::sparseMatrix( i = seq_len(Nobs), j = Adf$r1 ), r2 = Matrix::sparseMatrix( i = seq_len(Nobs), j = Adf$r2 ) ) Amat <- do.call(cbind, AmatList) beta <- rep(1, ncol(X)) thetaOrig <- c(0.1, 0.1, 0.2) # SD for each random group, then overdispersion theta <- log(thetaOrig) gamma <- rnorm( ncol(Amat), sd = rep(thetaOrig[1:2], c(Nrandom1, Nrandom2)) ) eta_true <- as.vector(X %*% beta + Amat %*% gamma) mu_true <- exp(eta_true) Z <- rgamma(length(mu_true), thetaOrig[3]^(-2), thetaOrig[3]^(-2)) y <- rpois(length(mu_true), mu_true * Z) config <- list( beta = beta, theta = log(thetaOrig), transform_theta = TRUE, gamma = rep(0, ncol(Amat)), shards = adlaplace::ad_shards(Amat, num_shards = 1000), verbose = FALSE, package = "adlaplace" ) n_beta <- length(config$beta) n_gamma <- length(config$gamma) n_theta <- length(config$theta) ``` ## Building AD shards Each model piece is compiled separately with `ad_fun_ptr()`. The full model has four shards: observation likelihood, hyperparameter (extra) terms, and one shard per random-effect group. Shards are merged later with `ad_fun()`. `beta_map` and `gamma_map` map local design columns of `X` and `A` into the global `(beta, gamma, theta)` vector: column `j` of `beta_map` has one structural `1` at row `r` when local `X` column `j` uses global beta `r` (and similarly for `gamma_map` / `A`). When `ncol(X) > 0`, `ncol(beta_map)` must equal `ncol(X)` with exactly one nonzero per column. ### Observations and hyperparameters The observation shard uses the third $\theta$ component (global overdispersion). The extra shard contributes the prior on $\theta$. ```{r adShards_obs} data_obs <- adlaplace:::ad_data( y = y, A = Amat, X = X, theta_map = Matrix::sparseMatrix( i = n_theta, j = 1L, dims = c(n_theta, 1L) ), ad_kind = "observations", ad_fun = "nbinom_obs" ) data_extra <- adlaplace::ad_data( data_obs, ad_kind = "parameters", ad_fun = "nbinom_extra" ) ad_fun_obs <- adlaplace::ad_fun_ptr(data = data_obs, config = config) ad_fun_extra <- adlaplace::ad_fun_ptr(data = data_extra, config = config) ``` ### Random effects Each group gets a diagonal normal prior with its own $\theta$ index. ```{r adShards_random} model_r1 <- adlaplace:::ad_data( beta_map = n_beta, gamma_map = Matrix::sparseMatrix( i = seq_len(Nrandom1), j = seq_len(Nrandom1), dims = c(n_gamma, Nrandom1) ), theta_map = c(1L, n_theta), ad_kind = "random", ad_fun = "random_diagonal", precision = rep(1, Nrandom1) ) ad_fun_r1 <- adlaplace::ad_fun_ptr(data = model_r1, config = config) model_r2 <- adlaplace:::ad_data( beta_map = n_beta, gamma_map = Matrix::sparseMatrix( i = seq.int(Nrandom1 + 1L, length.out = Nrandom2), j = seq_len(Nrandom2), dims = c(n_gamma, Nrandom2) ), theta_map = c(2L, n_theta), ad_kind = "random", ad_fun = "random_diagonal", precision = rep(1, Nrandom2) ) ad_fun_r2 <- adlaplace::ad_fun_ptr(data = model_r2, config = config) ``` ## Observation likelihood check Scan the overdispersion parameter while holding $(\beta, \gamma, \theta_{1:2})$ fixed at the data-generating values. The AD joint log density (observation + extra shards) should match `stats::dnbinom` on the full sample. ```{r obsLogDens} xx <- c(beta, gamma, config$theta) log_theta_true <- config$theta[n_theta] nb_log_dens <- function(log_theta) { r <- exp(-2 * log_theta) sum(stats::dnbinom(y, mu = mu_true, size = r, log = TRUE)) } ad_nb_log_dens <- function(log_theta) { xx_scan <- xx xx_scan[length(xx_scan)] <- log_theta adlaplace::joint_log_dens(ad_fun_obs, xx_scan, negative = FALSE) + adlaplace::joint_log_dens(ad_fun_extra, xx_scan, negative = FALSE) } seq_log_theta <- log_theta_true + seq(-0.5, 0.5, length.out = 7) cmp <- data.frame( theta = exp(seq_log_theta), stats = vapply(seq_log_theta, nb_log_dens, numeric(1)), ad = vapply(seq_log_theta, ad_nb_log_dens, numeric(1)) ) cmp$diff <- cmp$ad - cmp$stats cmp ``` ## Combining shards Shard handles are combined with `c()`. This uses move semantics: after `ad_fun()` is called, the individual `ad_fun_ptr` objects are cleared and must not be reused. ```{r adShards_combined} ad_fun_plain <- c( ad_fun_obs, ad_fun_extra, ad_fun_r1, ad_fun_r2 ) ad_pack <- adlaplace::ad_fun(ad_fun_plain, num_threads = num_threads) ``` The per-shard log densities sum to the joint log density at a fixed parameter vector. ```{r checkByShard} config$gamma <- rep(0.01, length(config$gamma)) x_full <- c(config$beta, config$gamma, config$theta) shards <- seq.int(from = 0L, length.out = adlaplace:::n_groups(ad_fun_plain)) by_shard <- vapply( shards, function(s) { adlaplace::joint_log_dens( ad_fun_plain, x_full, shards = s, negative = TRUE ) }, numeric(1) ) res_fdf <- fun_obj_fdfh( ad_pack, c(config$beta, config$theta), config$gamma, inner = FALSE, verbose = TRUE ) by_shard[seq(1, min(c(5, length(by_shard))))] c( sum(by_shard), adlaplace::joint_log_dens(ad_fun_plain, x_full, negative = TRUE), adlaplace::joint_log_dens(ad_pack, x_full, negative = TRUE), res_fdf$f ) ``` ## Joint log density derivatives Gradients and Hessians can be evaluated on all shards or on a subset. Results should agree when the subset is the full model. ```{r testLogDens} log_all <- adlaplace::joint_log_dens(ad_fun_plain, x_full, negative = TRUE) log_sub <- adlaplace::joint_log_dens( ad_fun_plain, x_full, shards = 0:3, negative = TRUE ) grad_all <- adlaplace::grad(ad_fun_plain, x_full, negative = TRUE) grad_sub <- adlaplace::grad( ad_fun_plain, x_full, shards = 0:3, negative = TRUE ) hess_all <- adlaplace::hessian(ad_fun_plain, x_full, negative = TRUE) hess_sub <- adlaplace::hessian( ad_fun_plain, x_full, shards = 0:3, negative = TRUE ) c( log_diff = log_all - res_fdf$f, max_abs_grad_diff = max(abs(grad_all - res_fdf$grad)), max_abs_hess_diff = max(abs( as.matrix(hess_all) - as.matrix(res_fdf$hessian) )) ) ``` ## Inner optimization Given outer parameters $(\beta, \theta)$, `inner_opt()` maximizes over $\gamma$ (equivalently minimizes $-\log p(y, \gamma \mid \beta, \theta)$). The returned solution should track the simulated random effects. ```{r testInnerOpt} x_outer <- c(config$beta, config$theta) x_full <- c(config$beta, config$gamma, config$theta) inner_res <- adlaplace::inner_opt( parameters = x_outer, gamma = config$gamma, ad_fun = ad_pack, control = list( maxit = 2, report.level = 3, report.freq = 1, grad.tol = 1e-8, cg.tol = 1e-6 ), deriv = TRUE, verbose = TRUE ) adlaplace::joint_log_dens(ad_pack, x_full, negative = TRUE) inner_res$opt$fval plot( inner_res$opt$solution, gamma, xlab = "estimated", ylab = "true" ) abline(0, 1, col = "red", lty = 2) ``` ## likelihood `log_lik_laplace()` runs the inner optimizer and returns the Laplace approximation to the marginal log likelihood. With `deriv = TRUE`, it also returns the profile gradient and third-order trace terms used in the determinant correction. OpenMP shard affinity is fixed when you call `ad_fun(ad_fun_plain, num_threads = …)` (`owner_thread = shard_index %% num_threads`); `log_lik_laplace()` does not change it. The same assignment is used for inner optimization, outer derivatives when `deriv = TRUE`, and the profile trace (`trace_hinv_t` logic, same OpenMP shard groups). `inner_opt()` uses one OpenMP/CppAD parallel session for trust-region inner optimization, outer `get_fdfh`, and trace evaluation when `deriv = TRUE`. `ad_fun_ptr()` and `joint_log_dens()` on a plain ptr do not use threads. For serial runs only, use `ad_fun(..., num_threads = 1L)`. For thread-affinity diagnostics, reinstall from a fresh tarball with `PKG_CXXFLAGS="-DDEBUG" R CMD INSTALL adlaplace_*.tar.gz` (use `R CMD INSTALL --clean` if installing from source). Configure prints when DEBUG is enabled. With `config$verbose = TRUE`, `inner_opt` prints CppAD session phase messages (trust region, outer `get_fdfh`, trace shards). With `deriv = TRUE`, trace runs inside that session; standalone `trace_hinv_t()` remains for direct use. Each `CppadParallelScope` end flushes per-thread CppAD pools before the next `.Call`. ```{r testLogLik} res_noderiv <- adlaplace::log_lik_laplace( x = x_outer, ad_fun = ad_pack, config = config, deriv = FALSE ) res_deriv <- adlaplace::log_lik_laplace( x = x_outer, ad_fun = ad_pack, config = config, deriv = TRUE ) data.frame( log_lik = c(res_noderiv$log_lik, res_deriv$log_lik), neg_log_lik = c(res_noderiv$neg_log_lik, res_deriv$neg_log_lik), inner_fval = c(res_deriv$opt$fval, res_deriv$opt$fval), row.names = c("deriv = FALSE", "deriv = TRUE") ) summary(res_deriv$extra$trace3) res_deriv$grad ``` ## Outer optimization `outer_fn()` and `outer_gr()` wrap `log_lik_laplace()` for use with `stats::optim()`. A small cache environment stores the current inner $\gamma$ start. ```{r trustOptimOuterWrappers} cache <- new.env(parent = emptyenv()) cache$gamma <- config$gamma control_inner <- list( maxit = 100, report.level = 0, report.freq = 0 ) optim_control <- list(maxit = 200, trace = 0) bounds <- list(lower = rep(-5, length(x_outer)), upper = rep(5, length(x_outer))) common_optim_args <- c( list( par = x_outer, fn = adlaplace::outer_fn, method = "L-BFGS-B", control = optim_control, config = config, ad_fun = ad_pack, cache = cache, control_inner = control_inner ), bounds ) outer_fit_no_g <- do.call(stats::optim, common_optim_args) outer_fit <- do.call(stats::optim, c(common_optim_args, list(gr = adlaplace::outer_gr))) fit_cmp <- rbind( optim = outer_fit$par, optim_no_grad = outer_fit_no_g$par, true = c(beta, log(thetaOrig)) ) colnames(fit_cmp) <- c( paste0("beta[", seq_len(n_beta), "]"), paste0("log_theta[", seq_len(n_theta), "]") ) fit_cmp exp(outer_fit$par[n_beta + seq_len(n_theta)]) outer_fit$value ``` ## Derivative checks Finite-difference checks along one outer coordinate validate the profile gradient, the determinant derivative, and the chain rule through the inner mode $u(\beta, \theta)$. ```{r testLogLgrad, fig.height=7} if (!requireNamespace("abind", quietly = TRUE)) { stop("Install suggested package 'abind' to run derivative checks.") } config$verbose <- FALSE par(mfrow = c(3, 2), mar = c(2, 2, 2, 0), mgp = c(1, 0.5, 0)) x1 <- outer_fit$par Dpar <- length(x1) Ngrid <- 7L par_grid <- matrix(x1, nrow = Ngrid, ncol = Dpar, byrow = TRUE) Sx <- x1[Dpar] + seq(-0.5, 0.5, length.out = Ngrid) SxD <- Sx[-1] - diff(Sx) / 2 par_grid[, Dpar] <- Sx res_scan <- lapply( split(par_grid, row(par_grid)), adlaplace::log_lik_laplace, ad_fun = ad_pack, config = config, deriv = TRUE, gamma = cache$gamma ) SnegLik <- vapply(res_scan, `[[`, numeric(1), "neg_log_lik") Sdet <- vapply(res_scan, function(r) r$extra$hessian$half_log_det, numeric(1)) grad_mat <- do.call(rbind, lapply(res_scan, `[[`, "grad")) extra_df <- do.call(abind::abind, c(lapply(res_scan, `[[`, "deriv"), along = 3)) dU <- do.call( abind::abind, c(lapply(res_scan, function(r) as.matrix(r$extra$dU)), along = 3) ) u_hat <- do.call(rbind, lapply(res_scan, function(r) r$opt$solution)) res_mid <- res_scan[[(Ngrid + 1L) %/% 2L]] plot(Sx, SnegLik, type = "l", xlab = expression(theta[3]), ylab = "-log lik") plot(Sx, grad_mat[, Dpar], type = "l", xlab = expression(theta[3]), ylab = "grad") points(SxD, diff(SnegLik) / diff(Sx), pch = 16) legend("topright", legend = c("AD", "finite diff"), lty = c(1, NA), pch = c(NA, 1), bty = "n") Du <- 2L plot(Sx, u_hat[, Du], type = "l", xlab = expression(theta[3]), ylab = expression(u[Du])) plot( Sx, dU[Du, Dpar, ], type = "l", xlab = expression(theta[3]), ylab = expression(d * u[Du] / d * theta[3]) ) points(SxD, diff(u_hat[, Du]) / diff(Sx), pch = 16) plot(Sx, Sdet, type = "l", xlab = expression(theta[3]), ylab = "half log det") plot( Sx, extra_df[Dpar, "d_det", ], type = "l", xlab = expression(theta[3]), ylab = expression(d / d * theta[3] ~ log ~ det) ) points(SxD, diff(Sdet) / diff(Sx), pch = 16) hessian_outer <- Matrix::forceSymmetric(res_mid$extra$hessian$outer) hessian_plain <- adlaplace::hessian( ad_fun_plain, res_mid$full_parameters, inner = FALSE, negative = TRUE ) max_abs_hess_diff <- max(abs((hessian_outer - hessian_plain)@x), na.rm = TRUE) max_abs_hess_diff ``` ```{r checkDu} seq_gamma1 <- seq.int( ad_pack@sizes["beta"] + 1L, length.out = ad_pack@sizes["gamma"] ) # dU <- -Hstuff$H_inv %*% hessian_pack$outer[seq_gamma1, -seq_gamma1] dU_check <- -solve( res_mid$extra$hessian$inner, res_mid$extra$hessian$outer[seq_gamma1, -seq_gamma1] ) dU_check2 <- - res_mid$extra$hessian$H_inv %*% res_mid$extra$hessian$outer[seq_gamma1, -seq_gamma1] quantile(abs(dU_check - res_mid$extra$dU)) quantile(abs(dU_check - dU_check2)) quantile(abs(dU_check2 - res_mid$extra$dU)) quantile(as.matrix(solve(res_mid$extra$hessian$inner)) - as.matrix(res_mid$extra$hessian$H_inv)) quantile(as.matrix(solve(res_mid$extra$hessian$inner)) - Matrix::tcrossprod(res_mid$extra$hessian$half_H_inv)) quantile(as.matrix(res_mid$extra$hessian$H_inv) - Matrix::tcrossprod(res_mid$extra$hessian$half_H_inv)) (as.matrix(solve(res_mid$extra$hessian$inner)) - as.matrix(res_mid$extra$hessian$H_inv))[1:5, 1:5] par(mfrow = c(1, 1)) plot( Sx, dU[Du, Dpar, ], type = "l", xlab = expression(theta[3]), ylab = expression(d * u[Du] / d * theta[3]) ) points(SxD, diff(u_hat[, Du]) / diff(Sx), pch = 16) points(rep(res_mid$parameters[Dpar], 2), c(dU_check[Du, Dpar], dU_check2[Du, Dpar]), col = c("red", "blue"), cex = c(5, 3) ) ``` ## Joint density derivative checks The same finite-difference idea applies to the full joint log density $-\log p(\beta, \gamma, \theta \mid y)$ and its AD gradient and Hessian. ```{r derivJointDens, fig.height=5} par(mfrow = c(4, 2), mar = c(2, 2, 2, 0), mgp = c(1, 0.5, 0)) x_here <- res_mid$full_parameters Dpar_dens <- length(x_here) - 1 Ngrid <- 11L shards <- seq.int(from = 0L, length.out = adlaplace:::n_groups(ad_fun_plain)) par_grid <- matrix(x_here, nrow = Ngrid, ncol = length(x_here), byrow = TRUE) Sx <- x_here[Dpar_dens] + seq(-0.1, 0.1, length.out = Ngrid) SxD <- Sx[-1] - diff(Sx) / 2 par_grid[, Dpar_dens] <- Sx x_list <- split(par_grid, row(par_grid)) dens <- vapply( x_list, adlaplace::joint_log_dens, numeric(1), ad_fun = ad_fun_plain, shards = shards, negative = FALSE ) rbind( dens[1:5], dens[seq(to = length(dens), length.out = 5)] ) grad <- do.call( cbind, lapply( x_list, adlaplace::grad, ad_fun = ad_fun_plain, inner = FALSE, shards = shards, negative = FALSE ) ) plot(Sx, grad[Dpar_dens, ], type = "l", ylab = "AD gradient") points(SxD, diff(dens) / diff(Sx), pch = 16) legend("topright", legend = c("AD", "finite diff"), lty = c(1, NA), pch = c(NA, 1), bty = "n") hes <- array( dim = c(length(x_here), length(x_here), Ngrid), dimnames = list(NULL, NULL, NULL) ) for (i in seq_len(Ngrid)) { hes[, , i] <- as.matrix(adlaplace::hessian( ad_fun_plain, x_list[[i]], inner = FALSE, shards = shards, negative = FALSE )) } grad_slope <- apply(grad, 1, diff) / mean(diff(Sx)) s_par2 <- sort(unique(c(Dpar_dens, 1:3, seq(to = length(x_here), length.out = 5)))) s_par2 <- s_par2[seq(1, min(c(prod(par("mfrow")) - 1, length(s_par2))))] for (Dpar2 in s_par2) { plot( Sx, hes[Dpar_dens, Dpar2, ], type = "l", ylab = paste0("H[", Dpar_dens, ",", Dpar2, "]"), ylim = range(c(hes[Dpar_dens, Dpar2, ], grad_slope[, Dpar2]), na.rm = TRUE) ) points(SxD, grad_slope[, Dpar2], pch = 16) } ``` ## nicer interface Code verifying the omnibus interface gives the same results as barebones implementation. ```{r checkNicerInterface} dat <- cbind( data.frame(y = y), Adf, x = X[, 2] ) md <- model_data( data = dat, formula = nbinom(y, lower = 1e-9, init = 0.15) ~ x + iid(r1, init = 0.1) + iid(r2, init = 0.1) ) config2 <- list( transform_theta = TRUE, shards = ad_shards( md$data$A, num_shards = 100 ) ) ad_fun2 <- ad_fun(md, config2, num_threads = num_threads) md$observations$y@XTp[, 1:5] data_obs@XTp[, 1:5] md$data$info$parameters[, c("term", "model", "label", "init", "transform")] data_obs@theta_map param_remap <- c(2, 1, 5, 3, 4) log_lik2 <- log_lik_laplace( x = res_deriv$parameters[param_remap], config = config2, ad_fun = ad_fun2, deriv = TRUE ) c(log_lik2$log_lik, res_deriv$log_lik) rbind(log_lik2$grad, res_deriv$grad[param_remap]) ``` # mean estimation Posterior mean of the random effects can be approximated with sparse nested Gauss–Hermite quadrature (`mvQuad`): nodes are drawn from the Laplace proposal `N(hat(gamma), H^{-1})` and reweighted by `joint_log_dens` relative to that Gaussian. ```{r meanEstimation} library(mvQuad) mean_estimate <- mean_mvquad( parameters = log_lik2$parameters, mode = log_lik2$opt$solution, cov = log_lik2$extra$hessian$H_inv, ad_fun = ad_fun2, n = 3 ) quantile(mean_estimate - log_lik2$opt$solution) plot(log_lik2$opt$solution, mean_estimate) abline(0, 1, col = "red", lty = 2) ```