---
title: "Laplace Approximations with adlaplaceExample"
author: "Patrick Brown"
date: "`r Sys.Date()`"
output:
  rmarkdown::html_vignette:
    toc: true
vignette: >
  %\VignetteIndexEntry{Laplace Approximations with adlaplaceExample}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

**Long-form lab notebook.** This vignette is intentionally detailed: it records
end-to-end development checks (simulation, shard construction, derivative
validation, inner Laplace solves, and outer optimization) rather than a minimal
getting-started tutorial. Skim the section headings if you only need a quick
overview; work through the chunks in order when reproducing or extending the
example backend.

This vignette walks through a hierarchical skew-normal model using the
`adlaplaceExample` backend: 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 skew-normal with location $\xi = \eta$, scale $\omega$, and
shape $\alpha$. The log density for one observation is

$$
\ell(y;\xi,\omega,\alpha)
= -\log\omega - \frac12\log(2\pi)
  -\frac12\left(\frac{y-\xi}{\omega}\right)^2
  + \log\left[\operatorname{erfc}\left(-\frac{\alpha(y-\xi)}{\sqrt{2}\,\omega}\right)\right].
$$

The linear predictor is $\eta = X\beta + A\gamma$, with two independent random-effect
groups in $A$, log-standard-deviation parameters for each group, plus skew-normal
scale $\omega$ (log scale when `transform_theta = TRUE`) and shape $\alpha$.
Densities are implemented in `src/objectiveFunction.cpp`; tapes for custom shards
are built via `adlaplace::ad_fun_ptr()` using the `package` slot on each `ad_data`
shard (`adlaplaceExample` for obs/extra, `adlaplace` for built-in random effects).

## Simulated data

We generate 5000 observations, fixed effects, two random-effect groups, and
skew-normal scale and shape. Group SDs and $\omega$ are stored on the log scale
(`transform_theta = TRUE`); $\alpha$ is untransformed.

```{r theData}
library(Matrix)
library(adlaplace)
library(adlaplaceExample)

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_len(c(1, 0.5), ncol(X))
thetaOrig <- c(sd1 = 0.1, sd2 = 0.1, omega = 0.25, alpha = 0.8)

gamma <- rnorm(
  ncol(Amat),
  sd = rep(thetaOrig[1:2], c(Nrandom1, Nrandom2))
)

eta_true <- as.vector(X %*% beta + Amat %*% gamma)

if (requireNamespace("sn", quietly = TRUE)) {
  y <- sn::rsn(Nobs, xi = eta_true, omega = thetaOrig["omega"], alpha = thetaOrig["alpha"])
} else {
  y <- stats::rnorm(Nobs, mean = eta_true, sd = thetaOrig["omega"])
}



config <- list(
  beta = beta,
  theta = c(log(thetaOrig[setdiff(names(thetaOrig), "alpha")]), thetaOrig["alpha"]),
  transform_theta = TRUE,
  gamma = rep(0, ncol(Amat)),
  shards = adlaplace::ad_shards(Amat, num_shards = 100),
  num_threads = 4L,
  verbose = FALSE
)

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()`.

### Observations and hyperparameters

The observation shard uses $\omega$ and $\alpha$ (last two $\theta$ components).
The extra shard contributes the $-\sum_i \log\omega$ normalization terms.

```{r adShards_obs}
data_obs <- adlaplace::ad_data(
  y = y,
  A = Amat,
  X = X,
  beta_map = Matrix::Diagonal(n_beta),
  gamma_map = Matrix::Diagonal(n_gamma),
  theta_map = list(c(
    which(names(thetaOrig) == "omega"),
    which(names(thetaOrig) == "alpha")
  ), n_theta),
  ad_kind = "observations",
  ad_fun = "skewnormal_obs",
  package = "adlaplaceExample"
)

data_extra <- adlaplace::ad_data(
  y = data_obs@y,
  beta_map = data_obs@beta_map,
  gamma_map = data_obs@gamma_map,
  theta_map = data_obs@theta_map,
  ad_kind = "parameters",
  ad_fun = "skewnormal_extra",
  package = "adlaplaceExample"
)
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}
nr1 <- Nrandom1
nr2 <- Nrandom2

model_r1 <- adlaplace::ad_data(
  beta_map = n_beta,
  gamma_map = Matrix::sparseMatrix(
    i = seq_len(nr1),
    j = seq_len(nr1),
    dims = c(n_gamma, nr1)
  ),
  theta_map = c(1L, n_theta),
  ad_kind = "random",
  ad_fun = "random_diagonal",
  package = "adlaplace",
  precision = rep(1, nr1)
)
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(nr1 + 1L, length.out = nr2),
    j = seq_len(nr2),
    dims = c(n_gamma, nr2)
  ),
  theta_map = c(2L, n_theta),
  ad_kind = "random",
  ad_fun = "random_diagonal",
  package = "adlaplace",
  precision = rep(1, nr2)
)
ad_fun_r2 <- adlaplace::ad_fun_ptr(data = model_r2, config = config)
```

## Observation likelihood check

Scan $\log(\omega)$ while holding $(\beta, \gamma, \theta_{1:2}, \alpha)$ fixed at
the data-generating values. The AD joint log density (observation + extra shards)
should match `sn::dsn` on the full sample when the **sn** package is available.

```{r obsLogDens}
if (!requireNamespace("sn", quietly = TRUE)) {
  stop("Install package 'sn' to run the observation likelihood check.")
}

xx <- c(beta, gamma, config$theta)

log_omega_true <- config$theta["omega"]
alpha_true <- config$theta["alpha"]
omega_idx <- n_beta + n_gamma + which(names(thetaOrig) == "omega")

eta_true <- as.vector(X %*% beta + Amat %*% gamma)

sn_log_dens <- function(log_omega) {
  omega <- exp(log_omega)
  sum(sn::dsn(y, xi = eta_true, omega = omega, alpha = alpha_true, log = TRUE))
}
ad_sn_log_dens <- function(log_omega) {
  xx_scan <- xx
  xx_scan[omega_idx] <- log_omega
  adlaplace::joint_log_dens(ad_fun_obs, xx_scan, negative = FALSE)
}
ad_sn_log_dens_extra <- function(log_omega) {
  xx_scan <- xx
  xx_scan[omega_idx] <- log_omega
  adlaplace::joint_log_dens(ad_fun_extra, xx_scan, negative = FALSE)
}

# skewnormal_obs shard (logDensObs): sum_i -z_i^2 + log(erfc(-t_i)), z = (y - eta)/(omega sqrt(2)),
# t = alpha * z; log(erfc(-t)) = log(2) + log Phi(alpha * (y - eta) / omega)
sn_obs_log_dens <- function(log_omega) {
  omega <- exp(log_omega)
  z <- (y - eta_true) / (omega * sqrt(2))
  tt <- alpha_true * z
  sum(
    -z^2 +
      log(pracma::erfc(-tt))
  )
}

# skewnormal_extra shard (logDensExtra): N * (-log omega - 1/2 log(2 pi))
sn_extra_log_dens <- function(log_omega) {
  Nobs * (-log_omega - 0.5 * log(2 * pi))
}

seq_log_omega <- log_omega_true + seq(-0.5, 0.5, length.out = 7)
cmp <- data.frame(
  omega = exp(seq_log_omega),
  ad_obs = vapply(seq_log_omega, ad_sn_log_dens, numeric(1)),
  obs_ref = vapply(seq_log_omega, sn_obs_log_dens, numeric(1)),
  ad_extra = vapply(seq_log_omega, ad_sn_log_dens_extra, numeric(1)),
  extra_ref = vapply(seq_log_omega, sn_extra_log_dens, numeric(1)),
  sn = vapply(seq_log_omega, sn_log_dens, numeric(1))
)
cmp$ad <- cmp$ad_obs + cmp$ad_extra
cmp$sum_ref <- cmp$obs_ref + cmp$extra_ref
# cmp$diff <- cmp$ad - cmp$sn
cmp
range(cmp$sn - cmp$ad)
```

## 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_r1, ad_fun_r2, ad_fun_extra)
ad_pack <- adlaplace::ad_fun(ad_fun_plain, num_threads = config$num_threads)
```

The per-shard log densities sum to the joint log density at a fixed parameter
vector.

```{r checkByShard}
x_full <- c(config$beta, config$gamma, config$theta)
shards <- seq.int(from = 0L, length.out = adlaplace:::n_groups(ad_pack@ptr))

by_shard <- vapply(
  shards,
  function(s) {
    adlaplace::joint_log_dens(
      ad_pack, x_full,
      shards = s, negative = FALSE
    )
  },
  numeric(1)
)
names(by_shard) <- c("obs", "extra", "r1", "r2")
by_shard
sum(by_shard)
adlaplace::joint_log_dens(ad_fun_plain, x_full, negative = FALSE)
```



## 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)

inner_res <- adlaplace::inner_opt(
  parameters = x_outer,
  gamma = config$gamma,
  ad_fun = ad_pack,
  control = list(
    maxit = 100,
    report.level = 0,
    report.freq = 0
  ),
  deriv = TRUE,
  verbose = FALSE
)

plot(
  inner_res$opt$solution, gamma,
  xlab = "estimated", ylab = "true",
  main = "Inner mode vs simulated gamma"
)
abline(0, 1, col = "red", lty = 2)
```

## Profile 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.

```{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")
)

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 <- rep(0, n_gamma)



control_inner <- list(
  maxit = 100,
  report.level = 0,
  report.freq = 0
)

optim_control <- list(maxit = 2, trace = 0)
bounds <- list(lower = c(-2, -2, -5, -5, -2, -2), upper = c(2, 2, 0, 0, 3, 2))



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)))



true_outer <- c(beta, log(thetaOrig[1:3]), thetaOrig[4])
fit_cmp <- rbind(
  optim = outer_fit$par,
  optim_no_grad = outer_fit_no_g$par,
  true = true_outer
)
colnames(fit_cmp) <- c(
  paste0("beta[", seq_len(n_beta), "]"),
  names(thetaOrig)
)
fit_cmp
```

## Profile 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.")
}

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

x1 <- outer_fit$par
Dpar <- 3 # length(x1) - 1
Ngrid <- 7L

par_grid <- matrix(x1, nrow = Ngrid, ncol = length(x1), 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(alpha), ylab = "-log lik")
plot(Sx, grad_mat[, Dpar], type = "l", xlab = expression(alpha), ylab = "grad")
points(SxD, diff(SnegLik) / diff(Sx), pch = 16)
legend("bottomright", 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(alpha), ylab = expression(u[Du]))
diff_here <- diff(u_hat[, Du]) / diff(Sx)
plot(
  Sx, dU[Du, Dpar, ],
  type = "l",
  ylim = range(c(dU[Du, Dpar, ], diff_here), na.rm = TRUE),
  xlab = expression(alpha),
  ylab = expression(d * u[Du] / d * alpha)
)
points(SxD, diff_here, pch = 16)

plot(Sx, Sdet, type = "l", xlab = expression(alpha), ylab = "half log det")
diff_here <- diff(Sdet) / diff(Sx)
plot(
  Sx, extra_df[Dpar, "d_det", ],
  type = "l",
  ylim = range(c(extra_df[Dpar, "d_det", ], diff_here), na.rm = TRUE),
  xlab = expression(alpha),
  ylab = expression(d / d * alpha ~ log ~ det)
)
points(SxD, diff_here, 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
```

## 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(2, 3), 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
)

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))
for (Dpar2 in c(Dpar_dens, 1L, 5, 6, length(x_here))) {
  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 the barebones
implementation.

```{r checkNicerInterface}
dat <- cbind(
  data.frame(y = y),
  Adf,
  x = X[, 2]
)


md <- adlaplace::model_data(
  data = dat,
  formula =
    skewnormal(y) ~
      x + adlaplace::iid(r1, init = 0.1) +
      adlaplace::iid(r2, init = 0.1)
)


md$data$info$theta[, c("label", "transform", "id")]
md$observations$y@theta_map
data_obs@theta_map
names(thetaOrig)

# order of parameters is different, previously alpha was last
# also intercept is now second

param_remap <- c(2, 1, 5, 6, 3, 4)
xx_remap <- res_deriv$parameters[param_remap]


config2 <- list(
  transform_theta = TRUE,
  shards = adlaplace::ad_shards(
    md$data$A,
    num_shards = 100
  ),
  beta = xx_remap[seq_len(n_beta)],
  theta = xx_remap[-seq_len(n_beta)],
  gamma = rep(0, n_gamma),
  verbose = TRUE
)




xx_all <- c(
  res_deriv$parameters[seq_len(n_beta)],
  rep(0, n_gamma),
  res_deriv$parameters[-seq_len(n_beta)]
)
xx_all2 <- c(
  xx_remap[seq_len(n_beta)],
  rep(0, n_gamma),
  xx_remap[-seq_len(n_beta)]
)

ad_extra <- adlaplace::ad_fun_ptr(data = data_extra, config = config)
ad_extra2 <- adlaplace::ad_fun_ptr(data = md$parameters$y_extra, config = config2)
ad_pack2 <- adlaplace::ad_fun(
  md,
  config2,
  num_threads = 1L
)

adlaplace::joint_log_dens(ad_extra, xx_all, negative = FALSE)
adlaplace::joint_log_dens(ad_extra2, xx_all2, negative = FALSE)
adlaplace::joint_log_dens(ad_pack2, xx_all2,
  shards = adlaplace:::n_groups(ad_pack2@ptr) - 1L,
  negative = FALSE
)


by_shard <- vapply(
  seq.int(from = 0L, length.out = adlaplace:::n_groups(ad_pack@ptr)),
  function(s) adlaplace::joint_log_dens(ad_pack, xx_all, shards = s, negative = FALSE),
  numeric(1)
)
by_shard2 <- vapply(
  seq.int(from = 0L, length.out = adlaplace:::n_groups(ad_pack2@ptr)),
  function(s) adlaplace::joint_log_dens(ad_pack2, xx_all2, shards = s, negative = FALSE),
  numeric(1)
)

rbind(by_shard, by_shard2)[, seq(to = length(by_shard), len = 6)]

c(sum(by_shard), sum(by_shard2))
adlaplace::joint_log_dens(ad_pack2, xx_all2, negative = FALSE)

log_lik1 <- adlaplace::log_lik_laplace(
  x = res_deriv$parameters,
  config = config,
  ad_fun = ad_pack,
  deriv = TRUE
)
log_lik2 <- adlaplace::log_lik_laplace(
  x = xx_remap,
  config = config2,
  ad_fun = ad_pack2,
  deriv = TRUE
)
c(log_lik2$log_lik, log_lik1$log_lik)
rbind(log_lik2$grad, log_lik1$grad[param_remap])
```