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

simple examples using gamm as alternative to gam
taken from help files for mgcv::gamm

```{r setup, include=FALSE}
library("Matrix")
library("adlaplace")
```

```{r gamm}
# Use packaged data for vignette builds (stable across renders).
# To regenerate locally:
#   set.seed(0)
#   dat <- mgcv::gamSim(6, n = 500)
#   dat$f2 <- sqrt(dat$f2)
#   dat$f2 <- dat$f2 - dat$f2[which.min(abs(dat$x2 - 0.5))]
#   dat$mu <- exp(-1 + 0 * dat$x1 + dat$f2 +
#     rnorm(nlevels(dat$fac), sd = 0.25)[as.integer(dat$fac)])
#   nbinom_sd <- 0.25
#   dat$Z <- rgamma(nrow(dat), nbinom_sd^(-2), nbinom_sd^(-2))
#   dat$y <- rpois(nrow(dat), lambda = dat$mu * dat$Z)
#   save(dat, file = "data/gamm.RData")
data("gamm", package = "adlaplace")
```


```{r collect_terms}
my_formula <-
  adlaplace::nbinom(y, lower = 1e-9, init = 0.15) ~
  x1 +
  adlaplace::iwp(
    x2,
    ref_value = 0.5,
    p = 2, knots = seq(0, 1, len = 21),
    init = 1
  ) +
  adlaplace::iid(fac, init = 0.25)
```

```{r model_data}
md <- model_data(
  data = dat,
  formula = my_formula
)

config <- list(
  transform_theta = TRUE,
  shards = adlaplace::ad_shards(
    md$data$A,
    num_shards = 100
  ),
  verbose = FALSE
)
```


```{r ad_fun_create}
# 2 threads when this build has OpenMP; otherwise serial.
num_threads <- if (adlaplace::has_openmp()) 2L else 1L
config$num_threads <- num_threads
control_inner <- list(maxit = 100L, report.level = 0, report.freq = 0)
ad_fun <- adlaplace::ad_fun(md, config, num_threads = num_threads)
```



```{r log_lik}
res <- adlaplace::log_lik_laplace(
  x = md$data$info$parameters$init,
  ad_fun = ad_fun,
  config = c(config, list(verbose = FALSE)),
  control = control_inner,
  deriv = TRUE
)
res$parameters
res$log_lik
res$grad
```


```{r trustOptimOuterWrappers}
md$data$info$parameters$lower[is.infinite(md$data$info$parameters$lower)] <- -5
md$data$info$parameters$upper[is.infinite(md$data$info$parameters$upper)] <- 5

cache <- new.env(parent = emptyenv())
cache$gamma <- rep(0, nrow(md$data$info$gamma)) # config$gamma not required; warm-start via cache

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

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

outer_fit <- stats::optim(
  par = md$data$info$parameters$init,
  fn = adlaplace::outer_fn,
  gr = adlaplace::outer_gr,
  method = "L-BFGS-B",
  control = list(
    maxit = 200,
    trace = 1,
    factr = 1e3,
    pgtol = 1e-12,
    parscale = md$data$info$parameters$parscale
  ),
  lower = md$data$info$parameters$lower,
  upper = md$data$info$parameters$upper,
  config = config,
  ad_fun = ad_fun,
  cache = cache,
  control_inner = control_inner
)

outer_fit$par
outer_fit$summary <- format_parameters(
  parameters = outer_fit$par, gamma = cache$gamma,
  info = ad_fun@info
)
outer_fit$summary$parameters
outer_fit$value
adlaplace::outer_gr(x = outer_fit$par, cache = cache, config = config, ad_fun = ad_fun)
```

## check by shard

```{r checkByShard}
shards <- seq.int(from = 0, length.out = adlaplace:::n_groups(ad_fun@ptr))
xx <- c(
  md$data$info$beta$init,
  rep(0, nrow(md$data$info$gamma)),
  adlaplace::apply_theta_log(md$data$info$theta, cols = "init")$init
)


by_shard <- mapply(
  adlaplace::joint_log_dens,
  shards = shards,
  MoreArgs = list(x = xx, ad_fun = ad_fun, negative = FALSE),
  SIMPLIFY = TRUE
)
by_shard[seq(to = length(by_shard), length.out = 5)]
```


## 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}
par(mfrow = c(3, 2), mar = c(2, 2, 2, 0), mgp = c(1, 0.5, 0))

x1 <- outer_fit$par
Dpar <- 4 # length(x1)
Ngrid <- 7L
n_beta <- nrow(md$data$info$beta)

Sx <- x1[Dpar] + 0.5 * seq(-1, 1, length.out = Ngrid)
SxD <- Sx[-1] - diff(Sx) / 2
par_grid <- matrix(x1,
  nrow = Ngrid, ncol = nrow(md$data$info$parameters),
  byrow = TRUE
)
par_grid[, Dpar] <- Sx

scan_label <- md$data$info$parameters$label[Dpar]

res_scan <- lapply(
  split(par_grid, row(par_grid)),
  adlaplace::log_lik_laplace,
  ad_fun = ad_fun,
  config = config,
  control = control_inner,
  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, xlab = scan_label, ylab = "-log lik")

plot(Sx, grad_mat[, Dpar], type = "l", xlab = scan_label, 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 <- min(2L, ncol(u_hat))
plot(Sx, u_hat[, Du], type = "l", xlab = scan_label, ylab = paste0("u[", Du, "]"))
plot(
  Sx, dU[Du, Dpar, ],
  type = "l",
  xlab = scan_label,
  ylab = paste0("du[", Du, "]/d", scan_label)
)
points(SxD, diff(u_hat[, Du]) / diff(Sx), pch = 16)

plot(Sx, Sdet, xlab = scan_label, ylab = "half log det")
plot(
  Sx, extra_df[Dpar, "d_det", ],
  type = "l",
  xlab = scan_label,
  ylab = paste0("d log det / d", scan_label)
)
points(SxD, diff(Sdet) / diff(Sx), pch = 16)

hessian_outer <- Matrix::forceSymmetric(res_mid$extra$hessian$outer)
hessian_ad <- adlaplace::hessian(
  ad_fun,
  res_mid$full_parameters,
  inner = FALSE,
  negative = TRUE
)
max(abs((hessian_outer - hessian_ad)@x), na.rm = TRUE)
```

## 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(3, 2), mar = c(2, 2, 2, 0), mgp = c(1, 0.5, 0))

x_here <- c(
  md$data$info$beta$init,
  0.1 + rep(0, nrow(md$data$info$gamma)),
  adlaplace::apply_theta_log(md$data$info$theta, cols = "init")$init
)
Dpar_dens <- 1 # length(x_here)
Ngrid <- 11L
shards <- seq.int(from = 0L, length.out = adlaplace:::n_groups(ad_fun@ptr))

par_grid <- matrix(x_here, nrow = Ngrid, ncol = length(x_here), byrow = TRUE)
Sx <- x_here[Dpar_dens] + 1 * seq(-1, 1, length.out = Ngrid)
SxD <- Sx[-1] - diff(Sx) / 2
par_grid[, Dpar_dens] <- Sx
x_list <- split(par_grid, row(par_grid))

dens <- mapply(
  adlaplace::joint_log_dens,
  x = x_list,
  MoreArgs = list(ad_fun = ad_fun, negative = FALSE)
)

grad <- do.call(
  cbind,
  lapply(
    x_list,
    adlaplace::grad,
    ad_fun = ad_fun,
    inner = FALSE,
    shards = shards,
    negative = FALSE
  )
)

plot(Sx, dens)

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,
    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, length(x_here) - 2L, length(x_here) - 1L)) {
  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)
}
```


```{r plot_results}
res <- adlaplace::log_lik_laplace(
  x = outer_fit$par,
  gamma = cache$gamma,
  ad_fun = ad_fun,
  config = config,
  control = control_inner,
  deriv = FALSE
)

newx <- data.frame(x2 = seq(0, 1, len = 101))

sim <- adlaplace::sim_fit(
  x = newx,
  data = md,
  fit = res,
  n = 500L
)

matplot(newx$x2, sim, type = "l", col = "#00000010", lty = 1)
points(dat$x2, dat$f2, col = "red", cex = 0.1)
```