SUN(3,3) maximum likelihood with analytic gradients

library(sn)
set.seed(0)
n <- 500
ysim <- sn::rsun(n, dp = true_sun_params)
colnames(ysim) <- paste0("Y", 1:3)
head(ysim)
#>              Y1         Y2          Y3
#> [1,]  2.0151802  0.6227806  1.74877173
#> [2,]  1.1496344  1.0112111 -0.56163779
#> [3,] -0.8262530 -0.2679811  0.09395353
#> [4,]  2.2321155  1.0156802 -0.36942726
#> [5,] -0.7340130 -0.2586762 -0.08056388
#> [6,] -0.1603508  0.4551714 -0.61204363

# Unused for d<=3 specialized CDF values; kept for API compatibility.
qmc_n_points <- 64L
qmc_n_shifts <- 2L

build_time <- system.time({
  tape <- dsun_fun(
    ysim,
    true_params,
    n_points = qmc_n_points,
    n_shifts = qmc_n_shifts,
    n_threads = 2L
  )
})["elapsed"]

check <- tape$eval(true_params, log = TRUE, deriv = 1L)
sn_ll <- sum(sn::dsun(ysim, dp = true_sun_params, log = TRUE))
c(
  sn = sn_ll,
  admvn = check$value,
  build_sec = unname(build_time),
  n_threads = tape$n_threads
)
#>        sn     admvn build_sec n_threads 
#> -2040.789 -2040.789     0.780     2.000

start_params <- make_sun33_start_from_normal(
  mu = colMeans(ysim),
  Sigma = stats::cov(ysim),
  skew_strength = 0.1
)

bnd <- sun33_bounds()
lower <- bnd$lower
upper <- bnd$upper
# Keep nu floors a bit above the package default for this demo.
lower[4:6] <- pmax(lower[4:6], 1e-3)
# Rescale free parameters so L-BFGS-B does not take wild Omega steps.
parscale <- pmax(abs(as.numeric(start_params)), 0.1)
optim_ctrl <- list(
  fnscale = -1,
  maxit = 200,
  parscale = parscale
)

obj <- tape$optim_fns()

fit_ad_time <- system.time({
  fit_ad <- optim(
    start_params,
    fn = obj$fn,
    gr = obj$gr,
    method = "L-BFGS-B",
    lower = lower,
    upper = upper,
    control = optim_ctrl
  )
})

fit_ad$convergence
#> [1] 0
fit_ad$value
#> [1] -2036.037
unname(fit_ad_time["elapsed"])
#> [1] 1.202

fit_trust_time <- system.time({
  fit_trust <- sun_mle(
    tape,
    start_params,
    control = list(maxit = 200L, report.freq = 0L, report.level = 0L)
  )
})
fit_trust$value
#> [1] -2033.638
fit_trust$iterations
#> [1] 200
unname(fit_trust_time["elapsed"])
#> [1] 27.844

sn_ll_at <- function(par) {
  sum(sn::dsun(ysim, dp = make_sun_params(par), log = TRUE))
}
ad_ll_at <- function(par) {
  tape$eval(par, log = TRUE, deriv = 0L)$value
}

compare <- data.frame(
  method = "ad_optim",
  sn_loglik = sn_ll_at(fit_ad$par),
  admvn_loglik = ad_ll_at(fit_ad$par),
  iterations = unname(fit_ad$counts[["gradient"]]),
  elapsed_sec = unname(fit_ad_time["elapsed"]),
  stringsAsFactors = FALSE
)
if (exists("fit_trust")) {
  compare <- rbind(
    data.frame(
      method = "ad_trustOptim",
      sn_loglik = sn_ll_at(fit_trust$par),
      admvn_loglik = ad_ll_at(fit_trust$par),
      iterations = as.numeric(fit_trust$iterations),
      elapsed_sec = unname(fit_trust_time["elapsed"]),
      stringsAsFactors = FALSE
    ),
    compare
  )
}
rownames(compare) <- NULL
compare
#>          method sn_loglik admvn_loglik iterations elapsed_sec
#> 1 ad_trustOptim -2033.639    -2033.638        200      27.844
#> 2      ad_optim -2036.037    -2036.037         22       1.202

par_hat <- as.numeric(if (exists("fit_trust")) fit_trust$par else fit_ad$par)
names(par_hat) <- names(true_params)
ll0 <- function(p) tape$eval(p, log = TRUE, deriv = 0L)$value
f0 <- ll0(par_hat)
se <- vapply(seq_along(par_hat), function(j) {
  h <- pmax(1e-5, 1e-4 * abs(par_hat[j]))
  p_up <- p_dn <- par_hat
  p_up[j] <- par_hat[j] + h
  p_dn[j] <- par_hat[j] - h
  if (j %in% 4:6) {
    # keep nu positive
    if (p_dn[j] <= 1e-8) p_dn[j] <- max(par_hat[j] * 0.5, 1e-6)
    h_eff2 <- (p_up[j] - par_hat[j]) * (par_hat[j] - p_dn[j])
    # for unequal steps use (f+ - f0)/(x+-x0) - (f0 - f-)/(x0-x-) over mean step
    g_up <- (ll0(p_up) - f0) / (p_up[j] - par_hat[j])
    g_dn <- (f0 - ll0(p_dn)) / (par_hat[j] - p_dn[j])
    Hjj <- (g_up - g_dn) / (0.5 * ((p_up[j] - par_hat[j]) + (par_hat[j] - p_dn[j])))
  } else {
    Hjj <- (ll0(p_up) - 2 * f0 + ll0(p_dn)) / h^2
  }
  if (is.finite(Hjj) && Hjj < 0) sqrt(-1 / Hjj) else NA_real_
}, numeric(1))
z_score <- (par_hat - as.numeric(true_params)) / se

recovery <- data.frame(
  true = as.numeric(true_params),
  start = as.numeric(start_params),
  mle = as.numeric(par_hat),
  se = se,
  z = z_score,
  row.names = names(true_params)
)
knitr::kable(recovery, digits = 3)

par_hat <- as.numeric(if (exists("fit_trust")) fit_trust$par else fit_ad$par)
par_names <- names(true_params)
if (is.null(par_names)) {
  par_names <- paste0("p", seq_along(par_hat))
}
names(par_hat) <- par_names

idx <- match(c("xi1", "nu1", "L11", "gamma12"), par_names)

ll <- function(p) tape$eval(p, log = TRUE, deriv = 0L)$value
ad_grad_j <- function(p, j) tape$eval(p, log = TRUE, deriv = 1L)$gradient[j]
fd_grad_j <- function(p, j, h) {
  p_up <- p
  p_dn <- p
  p_up[j] <- p[j] + h
  p_dn[j] <- p[j] - h
  (ll(p_up) - ll(p_dn)) / (2 * h)
}

op <- par(mfrow = c(2, 2), mar = c(3.2, 3.2, 2.2, 0.8), oma = c(0, 0, 1.5, 0))
on.exit(par(op), add = TRUE)

for (j in idx) {
  center <- par_hat[j]
  span <- pmax(0.15 * abs(center), 0.05)
  grid <- center + seq(-span, span, length.out = 9L)
  if (par_names[j] %in% c("nu1", "nu2", "nu3")) {
    grid <- grid[grid > 1e-3]
  }

  ad_vals <- numeric(length(grid))
  fd_vals <- numeric(length(grid))
  for (k in seq_along(grid)) {
    p <- par_hat
    p[j] <- grid[k]
    h <- pmax(1e-5, 1e-4 * abs(grid[k]))
    ad_vals[k] <- ad_grad_j(p, j)
    fd_vals[k] <- fd_grad_j(p, j, h)
  }

  ylim <- range(c(ad_vals, fd_vals), finite = TRUE)
  plot(
    grid, ad_vals,
    type = "l", lwd = 2, col = "#1b4f72",
    xlab = par_names[j], ylab = expression(partialdiff * ell / partialdiff * theta),
    main = par_names[j],
    ylim = ylim
  )
  lines(grid, fd_vals, lwd = 2, col = "#b03a2e", lty = 2)
  abline(v = center, col = "gray50", lty = 3)
}

legend(
  "topleft",
  legend = c("AD", "finite diff"),
  col = c("#1b4f72", "#b03a2e"),
  lty = c(1, 2), lwd = 2, bty = "n", cex = 0.8
)