AIS data in adlaplaceExample

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,
  shards = adlaplace::ad_shards(model_stuff$data$A, num_shards = 100),
  num_threads = 2L,
  verbose = FALSE
)

Laplace log-likelihood at starting values

ad_fun() 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_fun <- adlaplace::ad_fun(
  model_stuff,
  config = config,
  num_threads = config$num_threads
)
res <- adlaplace::log_lik_laplace(
  x = model_stuff$data$info$parameters$init,
  ad_fun = ad_fun,
  config = config,
  deriv = TRUE
)

## Beginning optimization
## 
## iter              f           nrm_gr                     status         radCG iter                  CG result
##   1   1688878.228723   85870.385201     Continuing - TR expand    2.000000      1         Intersect TR bound
##   2   1527402.149188   75622.182521     Continuing - TR expand    4.000000      1         Intersect TR bound
##   3   1265708.921210   55313.902431     Continuing - TR expand    8.000000      1         Intersect TR bound
##   4    983541.561603   16433.724581     Continuing - TR expand   16.000000      1         Intersect TR bound
##   5    953188.216184       0.338769                 Continuing   16.000000      6          Reached tolerance
##   6    953188.216171       0.000006                 Continuing   16.000000      6          Reached tolerance
##   7    953188.216171       0.000000                 Continuing   16.000000      6          Reached tolerance
## 
## Iteration has terminated
##   7    953188.216171       0.000000                    Success

res$parameters

## [1] 10.000000  0.000000 -2.302585  0.100000 -3.912023

res$neg_log_lik

## [1] 953221

res$grad

## [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$opt$solution

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

adlaplace::outer_fn(x = x0, cache = cache, config = config, ad_fun = ad_fun)

## [1] 953221

adlaplace::outer_gr(x = x0, cache = cache, config = config, ad_fun = ad_fun)

## [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$data$info$parameters$lower,
  upper = model_stuff$data$info$parameters$upper,
  method = "L-BFGS-B",
  control = list(
    maxit = 10,
    parscale = model_stuff$data$info$parameters$parscale,
    trace = 3,
    REPORT = 1
  ),
  config = config,
  ad_fun = ad_fun,
  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_fun@info)$parameters

##        term      model                label transform          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_fun = ad_fun,
  config = modifyList(config, list(verbose = FALSE)),
  deriv = TRUE
)

## Beginning optimization
## 
## iter          f        nrm_gr                     status         radCG iter                  CG result
##   1   501.055226    0.000072                 Continuing    1.000000      2          Reached tolerance
##   2   501.055226    0.000000                 Continuing    1.000000      2          Reached tolerance
## 
## Iteration has terminated
##   2   501.055226    0.000000                    Success

res$grad

## [1]   10.8998003  -15.7241785 -200.5131343   17.0562747   -0.6898023

model_stuff$data$info$gamma$mode <- res$opt$solution
model_stuff$data$info$beta$mle <- res$parameters[seq(1, nrow(model_stuff$data$info$beta))]

model_stuff$data$info$theta$mle <- res$parameters[-(1:nrow(model_stuff$data$info$beta))]
model_stuff$data$info$theta[
  model_stuff$data$info$theta$transform, "mle"
] <-
  exp(model_stuff$data$info$theta[model_stuff$data$info$theta$transform, "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$data$info$beta$mle[1],
      omega = model_stuff$data$info$theta$mle[1],
      alpha = model_stuff$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_fun = ad_fun,
  config = config,
  deriv = TRUE,
  gamma = cache$gamma
)

## Beginning optimization
## 
## iter           f         nrm_gr                     status         radCG 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
##   3   3180.471734     0.000012   Continuing - TR contract    0.500000      3          Reached tolerance
##   4   3180.471734     0.000012   Continuing - TR contract    0.250000      3          Reached tolerance
##   5   3180.471734     0.000012   Continuing - TR contract    0.125000      3          Reached tolerance
##   6   3180.471734     0.000012   Continuing - TR contract    0.062500      3          Reached tolerance
##   7   3180.471734     0.000012   Continuing - TR contract    0.031250      3          Reached tolerance
##   8   3180.471734     0.000012   Continuing - TR contract    0.015625      3          Reached tolerance
##   9   3180.471734     0.000012   Continuing - TR contract    0.007812      3          Reached tolerance
##  10   3180.471734     0.000012   Continuing - TR contract    0.003906      3          Reached tolerance
## 
## iter           f         nrm_gr                     status         radCG iter                  CG result
##  11   3180.471734     0.000012   Continuing - TR contract    0.001953      3          Reached tolerance
##  12   3180.471734     0.000012   Continuing - TR contract    0.000977      3          Reached tolerance
##  13   3180.471734     0.000012   Continuing - TR contract    0.000488      3          Reached tolerance
##  14   3180.471734     0.000012   Continuing - TR contract    0.000244      3          Reached tolerance
##  15   3180.471734     0.000012   Continuing - TR contract    0.000122      3          Reached tolerance
##  16   3180.471734     0.000012   Continuing - TR contract    0.000061      3          Reached tolerance
##  17   3180.471734     0.000012   Continuing - TR contract    0.000031      3          Reached tolerance
##  18   3180.471734     0.000012   Continuing - TR contract    0.000015      3          Reached tolerance
##  19   3180.471734     0.000012   Continuing - TR contract    0.000008      3          Reached tolerance
##  20   3180.471734     0.000012   Continuing - TR contract    0.000004      3          Reached tolerance
## 
## iter           f         nrm_gr                     status         radCG iter                  CG result
##  21   3180.471734     0.000012   Continuing - TR contract    0.000002      3          Reached tolerance
##  22   3180.471734     0.000012   Continuing - TR contract    0.000001      3          Reached tolerance
##  23   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
##  24   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
##  25   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
##  26   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
##  27   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
##  28   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
##  29   3180.471734     0.000012   Continuing - TR contract    0.000000      3          Reached tolerance
## 
## Iteration has terminated
##  29   3180.471734     0.000012Radius of trust region is less than stop.trust.radius
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         radCG 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
##   3   2400.869865     0.000000                 Continuing    1.000000      3          Reached tolerance
## 
## Iteration has terminated
##   3   2400.869865     0.000000                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         radCG 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
##   3   1741.045744     0.000006   Continuing - TR contract    0.500000      3          Reached tolerance
##   4   1741.045744     0.000006   Continuing - TR contract    0.250000      3          Reached tolerance
##   5   1741.045744     0.000006   Continuing - TR contract    0.125000      3          Reached tolerance
##   6   1741.045744     0.000006   Continuing - TR contract    0.062500      3          Reached tolerance
##   7   1741.045744     0.000006   Continuing - TR contract    0.031250      3          Reached tolerance
##   8   1741.045744     0.000006   Continuing - TR contract    0.015625      3          Reached tolerance
##   9   1741.045744     0.000006   Continuing - TR contract    0.007812      3          Reached tolerance
##  10   1741.045744     0.000006   Continuing - TR contract    0.003906      3          Reached tolerance
## 
## iter           f         nrm_gr                     status         radCG iter                  CG result
##  11   1741.045744     0.000006   Continuing - TR contract    0.001953      3          Reached tolerance
##  12   1741.045744     0.000006   Continuing - TR contract    0.000977      3          Reached tolerance
##  13   1741.045744     0.000006   Continuing - TR contract    0.000488      3          Reached tolerance
##  14   1741.045744     0.000006   Continuing - TR contract    0.000244      3          Reached tolerance
##  15   1741.045744     0.000006   Continuing - TR contract    0.000122      3          Reached tolerance
##  16   1741.045744     0.000006   Continuing - TR contract    0.000061      3          Reached tolerance
##  17   1741.045744     0.000006   Continuing - TR contract    0.000031      3          Reached tolerance
##  18   1741.045744     0.000006   Continuing - TR contract    0.000015      3          Reached tolerance
##  19   1741.045744     0.000006   Continuing - TR contract    0.000008      3          Reached tolerance
##  20   1741.045744     0.000006   Continuing - TR contract    0.000004      3          Reached tolerance
## 
## iter           f         nrm_gr                     status         radCG iter                  CG result
##  21   1741.045744     0.000006   Continuing - TR contract    0.000002      3          Reached tolerance
##  22   1741.045744     0.000006   Continuing - TR contract    0.000001      3          Reached tolerance
##  23   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
##  24   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
##  25   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
##  26   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
##  27   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
##  28   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
##  29   1741.045744     0.000006   Continuing - TR contract    0.000000      3          Reached tolerance
## 
## Iteration has terminated
##  29   1741.045744     0.000006Radius of trust region is less than stop.trust.radius
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         radCG 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         radCG iter                  CG result
##   1   820.004988    0.002208                 Continuing    1.000000      2          Reached tolerance
##   2   820.004988    0.000000                 Continuing    1.000000      2          Reached tolerance
## 
## Iteration has terminated
##   2   820.004988    0.000000                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         radCG iter                  CG result
##   1   577.250360    0.000107                 Continuing    1.000000      2          Reached tolerance
##   2   577.250360    0.000000                 Continuing    1.000000      2          Reached tolerance
## 
## Iteration has terminated
##   2   577.250360    0.000000                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         radCG iter                  CG result
##   1   501.055226    0.000000   Continuing - TR contract    0.500000      2          Reached tolerance
## 
## Iteration has terminated
##   1   501.055226    0.000000                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         radCG iter                  CG result
##   1   591.640043    0.000189                 Continuing    1.000000      2          Reached tolerance
##   2   591.640043    0.000000                 Continuing    1.000000      2          Reached tolerance
## 
## Iteration has terminated
##   2   591.640043    0.000000                    Success
## 
## Beginning optimization
## 
## iter          f        nrm_gr                     status         radCG iter                  CG result
##   1   807.825434    0.002720                 Continuing    1.000000      2          Reached tolerance
##   2   807.825434    0.000000                 Continuing    1.000000      2          Reached tolerance
## 
## Iteration has terminated
##   2   807.825434    0.000000                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         radCG 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         radCG 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
##   3   1578.686469     0.000000                 Continuing    1.000000      3          Reached tolerance
## 
## Iteration has terminated
##   3   1578.686469     0.000000                    Success
## 
## Beginning optimization
## 
## iter           f         nrm_gr                     status         radCG 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         radCG 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
##   3   2757.898167     0.000007   Continuing - TR contract    0.500000      3          Reached tolerance
##   4   2757.898167     0.000007   Continuing - TR contract    0.250000      3          Reached tolerance
##   5   2757.898167     0.000007   Continuing - TR contract    0.125000      3          Reached tolerance
##   6   2757.898167     0.000007   Continuing - TR contract    0.062500      3          Reached tolerance
##   7   2757.898167     0.000007   Continuing - TR contract    0.031250      3          Reached tolerance
##   8   2757.898167     0.000007   Continuing - TR contract    0.015625      3          Reached tolerance
##   9   2757.898167     0.000007   Continuing - TR contract    0.007812      3          Reached tolerance
##  10   2757.898167     0.000007   Continuing - TR contract    0.003906      3          Reached tolerance
## 
## iter           f         nrm_gr                     status         radCG iter                  CG result
##  11   2757.898167     0.000007   Continuing - TR contract    0.001953      3          Reached tolerance
##  12   2757.898167     0.000007   Continuing - TR contract    0.000977      3          Reached tolerance
##  13   2757.898167     0.000007   Continuing - TR contract    0.000488      3          Reached tolerance
##  14   2757.898167     0.000007   Continuing - TR contract    0.000244      3          Reached tolerance
##  15   2757.898167     0.000007   Continuing - TR contract    0.000122      3          Reached tolerance
##  16   2757.898167     0.000007   Continuing - TR contract    0.000061      3          Reached tolerance
##  17   2757.898167     0.000007   Continuing - TR contract    0.000031      3          Reached tolerance
##  18   2757.898167     0.000007   Continuing - TR contract    0.000015      3          Reached tolerance
##  19   2757.898167     0.000007   Continuing - TR contract    0.000008      3          Reached tolerance
##  20   2757.898167     0.000007   Continuing - TR contract    0.000004      3          Reached tolerance
## 
## iter           f         nrm_gr                     status         radCG iter                  CG result
##  21   2757.898167     0.000007   Continuing - TR contract    0.000002      3          Reached tolerance
##  22   2757.898167     0.000007   Continuing - TR contract    0.000001      3          Reached tolerance
##  23   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
##  24   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
##  25   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
##  26   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
##  27   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
##  28   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
##  29   2757.898167     0.000007   Continuing - TR contract    0.000000      3          Reached tolerance
## 
## Iteration has terminated
##  29   2757.898167     0.000007Radius of trust region is less than stop.trust.radius

SnegLik <- vapply(res, `[[`, numeric(1), "neg_log_lik")
Sdet <- vapply(res, function(r) r$extra$hessian$half_log_det, numeric(1))
grad_mat <- do.call(rbind, lapply(res, `[[`, "grad"))
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$extra$dU)), along = 3)
)
u_hat <- do.call(rbind, lapply(res, function(r) r$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))