AIS data in adlaplaceExample

Patrick Brown

2026-10-11

This vignette fits a skew-normal location model to body-mass index (BMI) from the Australian Institute of Sport (ais) dataset. We use the adlaplaceExample backend for the observation density and adlaplace for Laplace approximation, inner optimization, and outer profiling.

Data

The sn package provides the original ais data; if it is not installed we build a small stand-in so the vignette still knits.

if (requireNamespace("sn", quietly = TRUE)) {
  data("ais", package = "sn")
  hist(ais$BMI)
} else {
  ais <- data.frame(
    sex = rep(c("male", "female"), 5),
    sport = rep(1:5, each = 2),
    BMI = rnorm(10)
  )
}

Model and AD setup

We specify a skew-normal response for BMI with an intercept, fixed effect of sex, and an independent random effect for sport. Observation shards are split across 100 groups for parallel derivative evaluation.

model_stuff <- adlaplace::model_data(
  adlaplaceExample::skewnormal(BMI) ~
    adlaplace::intercept(init = 10, parscale = 10) +
    sex + adlaplace::iid(sport),
  data = ais,
  verbose = TRUE
)
## Model term adlaplaceExample::sk... parsed from formula
## Model term adlaplace::intercept... parsed from formula
## Model term sex... parsed from formula
## Model term adlaplace::iid(sport... parsed from formula
## variables:
## [1] "BMI"   "sex"   "sport"
## data has  202  rows
## Model term adlaplaceExample::sk... parsed from formula
## Model term adlaplace::intercept... parsed from formula
## Model term sex... parsed from formula
## Model term adlaplace::iid(sport... parsed from formula
config <- list(
  transform_theta = TRUE,
  obs_groups = adlaplace::obs_groups(model_stuff$term_data$A, num_shards = 100),
  num_threads = 2L,
  verbose = FALSE
)

Laplace log-likelihood at starting values

ad_pack() composes the observation, parameter, and random-effect shards into a single callable object. log_lik_laplace() solves the inner problem for random effects and returns the profile log-likelihood, gradient, and fitted parameters at the initial outer vector.

ad_pack <- adlaplace::ad_pack(
  model_stuff,
  config = config,
  num_threads = config$num_threads
)
res <- adlaplace::log_lik_laplace(
  x = model_stuff$term_data$info$parameters$init,
  ad_pack = ad_pack,
  config = config,
  deriv = TRUE
)
## Beginning optimization
## 
## iter              f           nrm_gr                     status         rad  CG iter                  CG result
##   1   1688878.228723   85870.385201     Continuing - TR expand    3.000000        1         Intersect TR bound
##   2   1454362.627571   70513.142184     Continuing - TR expand    9.000000        1         Intersect TR bound
##   3   1025463.404412   25781.228570     Continuing - TR expand   27.000000        1         Intersect TR bound
##   4    953188.216207       0.586849                 Continuing   27.000000        6          Reached tolerance
##   5    953188.216171       0.000010                 Continuing   27.000000        6          Reached tolerance
## 
## Iteration has terminated
##   5    953188.216171       0.000010                    Success
res$parameters
## [1] 10.000000  0.000000 -2.302585  0.100000 -3.912023
res$neg_log_lik
## [1] 953221
res$deriv$d_neg_log_lik
## [1] -1.370136e+05 -7.840007e+04 -1.067354e+06 -1.143206e+01 -8.397178e+05

Outer optimization

The outer objective profiles out random effects via a cache. We evaluate the objective and gradient once at the start, then run a short L-BFGS-B fit.

cache <- new.env(parent = emptyenv())
cache$gamma <- res$inner_opt$solution

x0 <- model_stuff$term_data$info$parameters$init

adlaplace::outer_fn(x = x0, cache = cache, config = config, ad_pack = ad_pack)
## [1] 953221
adlaplace::outer_gr(x = x0, cache = cache, config = config, ad_pack = ad_pack)
## [1] -1.370136e+05 -7.840007e+04 -1.067354e+06 -1.143206e+01 -8.397178e+05
outer_fit <- stats::optim(
  par = x0,
  fn = adlaplace::outer_fn,
  gr = adlaplace::outer_gr,
  lower = model_stuff$term_data$info$parameters$lower,
  upper = model_stuff$term_data$info$parameters$upper,
  method = "L-BFGS-B",
  control = list(
    maxit = 10,
    parscale = model_stuff$term_data$info$parameters$parscale,
    trace = 3,
    REPORT = 1
  ),
  config = config,
  ad_pack = ad_pack,
  cache = cache
)
## N = 5, M = 5 machine precision = 2.22045e-16
## At X0, 0 variables are exactly at the bounds
## At iterate     0  f=   9.5322e+05  |proj g|=   1.3701e+06
## At iterate     1  f =        88432  |proj g|=    2.2605e+05
## At iterate     2  f =        53785  |proj g|=    1.5368e+05
## At iterate     3  f =        18431  |proj g|=         40307
## At iterate     4  f =        13382  |proj g|=         24186
## At iterate     5  f =        10108  |proj g|=         18531
## At iterate     6  f =       5034.7  |proj g|=        9368.8
## At iterate     7  f =       2436.2  |proj g|=        4399.6
## At iterate     8  f =       1342.5  |proj g|=        2159.7
## At iterate     9  f =       830.06  |proj g|=        1023.4
## At iterate    10  f =          612  |proj g|=        472.62
## At iterate    11  f =       527.28  |proj g|=        200.51
## final  value 527.279450 
## stopped after 11 iterations
outer_fit$par
## [1] 23.060589934  0.230926775  0.693340643 -0.002664903 -3.539238360
outer_fit$value
## [1] 527.2795
adlaplace::format_parameters(
  parameters = outer_fit$par, gamma = cache$gamma,
  info = ad_pack@info
)$parameters
##        term      model                label   log          mle
## 1 intercept  intercept            intercept FALSE 23.060589934
## 2       sex     linear           sex_linear FALSE  0.230926775
## 3       BMI skewnormal BMI_skewnormal_omega  TRUE  2.000386963
## 4       BMI skewnormal BMI_skewnormal_alpha FALSE -0.002664903
## 5     sport        iid            sport_iid  TRUE  0.029035433

Fitted values and density overlay

After refitting at the outer MLE, we store updated fixed effects, variance components, and random-effect modes, then overlay the fitted skew-normal density on the BMI histogram (when sn is available).

res <- adlaplace::log_lik_laplace(
  x = outer_fit$par,
  ad_pack = ad_pack,
  config = modifyList(config, list(verbose = FALSE)),
  deriv = TRUE
)
## Beginning optimization
## 
## iter          f        nrm_gr                     status         rad  CG iter                  CG result
##   1   501.055226    0.000072                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   1   501.055226    0.000072                    Success
res$deriv$d_neg_log_lik
## [1]   10.8998003  -15.7241785 -200.5131343   17.0562747   -0.6898023
model_stuff$term_data$info$gamma$mode <- res$inner_opt$solution
model_stuff$term_data$info$beta$mle <- res$parameters[seq(1, nrow(model_stuff$term_data$info$beta))]

model_stuff$term_data$info$theta$mle <- res$parameters[-(1:nrow(model_stuff$term_data$info$beta))]
model_stuff$term_data$info$theta[
  model_stuff$term_data$info$theta$log, "mle"
] <-
  exp(model_stuff$term_data$info$theta[model_stuff$term_data$info$theta$log, "mle"])
hist(ais$BMI, prob = TRUE, breaks = 31)
x_seq <- seq(par("usr")[1], par("usr")[2], len = 1001)
if (requireNamespace("sn", quietly = TRUE)) {
  lines(x_seq,
    sn::dsn(x_seq,
      xi = model_stuff$term_data$info$beta$mle[1],
      omega = model_stuff$term_data$info$theta$mle[1],
      alpha = model_stuff$term_data$info$theta$mle[2]
    ),
    col = "red"
  )
}

Derivative check along one parameter

Finally we scan the fourth outer parameter and compare the analytic gradient to finite differences, along with inner-mode and Hessian-determinant sensitivities. This requires the suggested abind package.

if (!requireNamespace("abind", quietly = TRUE)) {
  stop("Install suggested package 'abind' to run derivative checks.")
}

x <- cache$last_par_fn # outer_fit$par
Dpar <- 4
Ngrid <- 13L

par_grid <- matrix(x, nrow = Ngrid, ncol = length(x), byrow = TRUE)
Sx <- 5 * seq(-1, 1, len = Ngrid) + x[Dpar]
SxD <- Sx[-1] - diff(Sx) / 2
par_grid[, Dpar] <- Sx

res <- lapply(
  split(par_grid, row(par_grid)),
  adlaplace::log_lik_laplace,
  ad_pack = ad_pack,
  config = config,
  deriv = TRUE,
  gamma = cache$gamma
)
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   3180.472717     1.593348                 Continuing    1.000000        3          Reached tolerance
##   2   3180.471734     0.000012                 Continuing    1.000000        3          Reached tolerance
## 
## Iteration has terminated
##   2   3180.471734     0.000012                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   2400.869969     0.513170                 Continuing    1.000000        3          Reached tolerance
##   2   2400.869865     0.000049                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   2   2400.869865     0.000049                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   1741.045751     0.125032                 Continuing    1.000000        3          Reached tolerance
##   2   1741.045744     0.000006                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   2   1741.045744     0.000006                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   1211.124332     0.022028                 Continuing    1.000000        2          Reached tolerance
##   2   1211.124332     0.000002                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   2   1211.124332     0.000002                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         rad  CG iter                  CG result
##   1   820.004988    0.002208                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   1   820.004988    0.002208                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         rad  CG iter                  CG result
##   1   577.250360    0.000107                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   1   577.250360    0.000107                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         rad  CG iter                  CG result
##   1   501.055226    0.000000                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   1   501.055226    0.000000                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         rad  CG iter                  CG result
##   1   591.640043    0.000189                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   1   591.640043    0.000189                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         rad  CG iter                  CG result
##   1   807.825434    0.002720                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   1   807.825434    0.002720                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   1138.703495     0.022864                 Continuing    1.000000        2          Reached tolerance
##   2   1138.703495     0.000001                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   2   1138.703495     0.000001                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   1578.686473     0.093941                 Continuing    1.000000        3          Reached tolerance
##   2   1578.686469     0.000007                 Continuing    1.000000        2          Reached tolerance
## 
## Iteration has terminated
##   2   1578.686469     0.000007                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   2121.079324     0.332690                 Continuing    1.000000        3          Reached tolerance
##   2   2121.079281     0.000001                 Continuing    1.000000        3          Reached tolerance
## 
## Iteration has terminated
##   2   2121.079281     0.000001                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         rad  CG iter                  CG result
##   1   2757.898501     0.943724                 Continuing    1.000000        3          Reached tolerance
##   2   2757.898167     0.000007                 Continuing    1.000000        3          Reached tolerance
## 
## Iteration has terminated
##   2   2757.898167     0.000007                    Success
SnegLik <- vapply(res, `[[`, numeric(1), "neg_log_lik")
Sdet <- vapply(res, function(r) r$hessian$half_log_det, numeric(1))
grad_mat <- do.call(rbind, lapply(res, function(r) r$deriv$d_neg_log_lik))
extra_df <- do.call(abind::abind, c(lapply(res, `[[`, "deriv"), along = 3))
dU <- do.call(
  abind::abind,
  c(lapply(res, function(r) as.matrix(r$gradient$outer$dU)), along = 3)
)
u_hat <- do.call(rbind, lapply(res, function(r) r$inner_opt$solution))

par(mfrow = c(3, 2), mar = c(2, 2, 2, 0), mgp = c(1, 0.5, 0))
plot(Sx, SnegLik)

plot(Sx, grad_mat[, Dpar], type = "l")
points(SxD, diff(SnegLik) / diff(Sx))

Du <- 1L
plot(Sx, u_hat[, Du])

plot(Sx, dU[Du, Dpar, ])
lines(SxD, diff(u_hat[, Du]) / diff(Sx))

plot(Sx, Sdet)

plot(Sx, extra_df[Dpar, "d_det", ])
lines(SxD, diff(Sdet) / diff(Sx))