SUN(3,3) MLE with joint hyperspherical parameterization

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,]  1.7883370  0.4081508  1.59061403
#> [2,]  1.1074220  0.9592382 -0.70678666
#> [3,] -0.8842272 -0.3764776 -0.02136006
#> [4,]  2.2899684  1.1767293 -0.26603871
#> [5,] -0.8308968 -0.3323442 -0.30570815
#> [6,] -0.2338628  0.3750492 -0.74566247

# 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_hs_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 
#> -2068.704 -2068.704     0.651     2.000

start_params <- make_sun33_hs_start_from_normal(
  mu = colMeans(ysim),
  Sigma = stats::cov(ysim),
  skew_strength = 0.15
)

bnd <- sun33_hs_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)
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] -2063.988
unname(fit_ad_time["elapsed"])
#> [1] 1.757

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] -2061.181
fit_trust$iterations
#> [1] 200
unname(fit_trust_time["elapsed"])
#> [1] 25.178

sn_ll_at <- function(par) {
  sum(sn::dsun(ysim, dp = make_sun_hs_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 -2061.182    -2061.181        200      25.178
#> 2      ad_optim -2063.988    -2063.988         32       1.757

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) {
    if (p_dn[j] <= 1e-8) p_dn[j] <- max(par_hat[j] * 0.5, 1e-6)
    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", "z41", "z21"), 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
)