SUN(3,3) KL matching and recovery

This vignette builds a known SUN\((3,3)\), forms a weighted discrete target from it, and minimises the KL / cross-entropy with fit_sun33_quad(). The fitting score uses the dsun_fun() tape (same as evaluate_sun33_slice()). The free parameters use the residual-correlation construction so that \(\operatorname{Var}(V)=\Sigma_{vv}\) is a correlation matrix (diag(Gamma) = 1) without an \(\omega_v\) rescale. Skew loadings \(L\) act on the correlation \(C_u\) of the standardized latent (scale \(D_u=\operatorname{diag}(\nu)\) with \(\nu_j\) the marginal standard deviations is applied outside), so \(|L_{ij}|\le 1\) is the natural regime.

Because several block-Cholesky coordinates can represent nearly the same SUN law, raw recovery of all 21 free parameters is often weak (especially the off-diagonal contrasts of \(L\)). We therefore check that the KL drops toward the value at the true parameters, that the fitted density matches on a hold-out sample, and that the sn-style dp summaries are close.

True SUN(3,3)

par_true <- c(
  xi1 = 0.2, xi2 = -0.1, xi3 = 0.05,
  nu1 = 1.0, nu2 = 0.8, nu3 = 1.2,
  omega12 = 0.15, omega13 = 0.05, omega23 = 0.10,
  L11 = 0.5, L12 = 0, L13 = 0,
  L21 = 0, L22 = 0.45, L23 = 0,
  L31 = 0, L32 = 0, L33 = 0.4,
  gamma12 = 0.2, gamma13 = 0.0, gamma23 = 0.1
)
dp_true <- make_sun_params(par_true)
str(dp_true, max.level = 1)
#> List of 5
#>  $ xi   : num [1:3] 0.2 -0.1 0.05
#>  $ Omega: num [1:3, 1:3] 1 0.119 0.06 0.119 0.64 ...
#>  $ Delta: num [1:3, 1:3] 0.5 0.0744 0.025 0.067 0.45 ...
#>  $ tau  : num [1:3] 0 0 0
#>  $ Gamma: num [1:3, 1:3] 1 0.18615 0.00999 0.18615 1 ...
diag(dp_true$Gamma)
#> [1] 1 1 1

Discrete target from the true law

Sample from the true SUN and use the empirical measure as \(p\), with logp from sn::dsun at the true dp.

library(sn)
set.seed(1)
n <- 600L
x <- sn::rsun(n, dp = dp_true)
w <- rep(1 / n, n)
logp_true <- as.numeric(sn::dsun(x, dp = dp_true, log = TRUE))
quad <- list(x = x, w = w, logp = logp_true)

kl_at_true <- evaluate_sun33_slice(par_true, quad)
kl_at_true
#>            kl       neg_Elq 
#> -4.199885e-10  4.082711e+00

Start far from the truth

Start from a nearly Gaussian SUN matched to the sample moments.

par_start <- make_sun33_start_from_normal(
  mu = colMeans(x),
  Sigma = stats::cov(x),
  skew_strength = 0.05
)
kl_at_start <- evaluate_sun33_slice(par_start, quad)
kl_at_start
#>           kl      neg_Elq 
#> -0.002330011  4.080381025

sel <- c("xi1", "xi2", "xi3", "nu1", "nu2", "nu3", "L11", "L22", "L33", "L12", "L21", "gamma12", "gamma13", "gamma23")
rbind(true = par_true[sel], start = par_start[sel])
#>             xi1        xi2       xi3       nu1       nu2      nu3  L11  L22
#> true  0.2000000 -0.1000000 0.0500000 1.0000000 0.8000000 1.200000 0.50 0.45
#> start 0.6161526  0.1986393 0.4415031 0.9502412 0.7580858 1.181817 0.05 0.05
#>        L33 L12 L21 gamma12 gamma13 gamma23
#> true  0.40   0   0     0.2       0     0.1
#> start 0.05   0   0     0.0       0     0.0

Minimise KL

fit <- fit_sun33_quad(
  quad = quad,
  start = par_start,
  mc.cores = 1L,
  obj_scale = 1e4,
  mvquad_opts = list(n_points = 64L, n_shifts = 2L),
  optim_opts = list(maxit = 250L, factr = 1e3, pgtol = 1e-8)
)
print_sun_kl_fit(fit, label = "SUN(3,3) KL match")
#> SUN(3,3) KL match
#>   engine=admvn  convergence=1  counts={function=272, gradient=272}
#>   message: NEW_X 
#>   neg_Elq: start=4.08038  fit=4.07299
#>   kl:      start=-0.00233001  fit=-0.00972204
#>   ||grad||_inf=20.0243  (objective scale, obj_scale=10000)
#>   gradient:
#>  [1]  -2.8079839  -4.4862290  -3.6131465   9.4625677   4.8186280  -0.8968356
#>  [7]  -0.8399226  -7.3524638   3.6929307  20.0243290   2.4568877  14.1532639
#> [13]  -2.8389189  14.0108841   4.1281993  -0.1138700 -10.5218928  13.7915130
#> [19]  -8.2186871   9.4262297  -0.6488676

Recovery diagnostics

par_hat <- fit$par
names(par_hat) <- names(par_true)
dp_hat <- fit$dp
kl_at_hat <- evaluate_sun33_slice(par_hat, quad)

data.frame(
  kl_true = unname(kl_at_true[["kl"]]),
  kl_start = unname(kl_at_start[["kl"]]),
  kl_hat = unname(kl_at_hat[["kl"]])
)
#>         kl_true     kl_start       kl_hat
#> 1 -4.199885e-10 -0.002330011 -0.009722039

# Hold-out mean log-density under true vs fitted dp
set.seed(2)
x_hold <- sn::rsun(400L, dp = dp_true)
ll_true <- mean(sn::dsun(x_hold, dp = dp_true, log = TRUE))
ll_hat <- mean(sn::dsun(x_hold, dp = dp_hat, log = TRUE))
ll_start <- mean(sn::dsun(x_hold, dp = make_sun_params(par_start), log = TRUE))
c(holdout_ll_true = ll_true, holdout_ll_start = ll_start, holdout_ll_hat = ll_hat)
#>  holdout_ll_true holdout_ll_start   holdout_ll_hat 
#>        -4.010964        -4.016123        -4.028659

cmp <- data.frame(
  true = par_true,
  start = as.numeric(par_start),
  hat = as.numeric(par_hat),
  abs_err = abs(as.numeric(par_hat) - par_true)
)
print(cmp[sel, ], digits = 4)
#>          true  start       hat  abs_err
#> xi1      0.20 0.6162  0.229065 0.029065
#> xi2     -0.10 0.1986  0.216844 0.316844
#> xi3      0.05 0.4415 -0.751589 0.801589
#> nu1      1.00 0.9502  1.141387 0.141387
#> nu2      0.80 0.7581  0.879031 0.079031
#> nu3      1.20 1.1818  1.474506 0.274506
#> L11      0.50 0.0500  0.502515 0.002515
#> L22      0.45 0.0500 -0.426340 0.876340
#> L33      0.40 0.0500  0.663031 0.263031
#> L12      0.00 0.0000  0.305214 0.305214
#> L21      0.00 0.0000  0.228078 0.228078
#> gamma12  0.20 0.0000  0.005760 0.194240
#> gamma13  0.00 0.0000 -0.029174 0.029174
#> gamma23  0.10 0.0000 -0.009855 0.109855

list(
  xi = rbind(true = dp_true$xi, hat = dp_hat$xi),
  Omega_diag = rbind(true = diag(dp_true$Omega), hat = diag(dp_hat$Omega)),
  Delta_diag = rbind(true = diag(dp_true$Delta), hat = diag(dp_hat$Delta)),
  max_abs_Omega = max(abs(dp_hat$Omega - dp_true$Omega)),
  max_abs_Delta = max(abs(dp_hat$Delta - dp_true$Delta)),
  max_abs_Gamma = max(abs(dp_hat$Gamma - dp_true$Gamma)),
  max_abs_par = max(abs(par_hat - par_true))
)
#> $xi
#>           [,1]       [,2]       [,3]
#> true 0.2000000 -0.1000000  0.0500000
#> hat  0.2290647  0.2168436 -0.7515888
#> 
#> $Omega_diag
#>          [,1]     [,2]     [,3]
#> true 1.000000 0.640000 1.440000
#> hat  1.302765 0.772695 2.174167
#> 
#> $Delta_diag
#>           [,1]       [,2]      [,3]
#> true 0.5000000  0.4500000 0.4000000
#> hat  0.6048751 -0.3517748 0.6378247
#> 
#> $max_abs_Omega
#> [1] 0.7341671
#> 
#> $max_abs_Delta
#> [1] 0.8017748
#> 
#> $max_abs_Gamma
#> [1] 0.2519988
#> 
#> $max_abs_par
#> [1] 0.8763403

par(mfrow = c(1, 2), mar = c(4, 4, 2, 1))

logq_hat <- as.numeric(sn::dsun(x, dp = dp_hat, log = TRUE))
plot(
  logp_true, logq_hat,
  xlab = "true log density", ylab = "fitted log density",
  main = "Density match on fitting sample"
)
abline(0, 1, col = "red", lty = 2)

plot(
  as.vector(dp_true$Omega), as.vector(dp_hat$Omega),
  xlab = "true Omega", ylab = "fitted Omega",
  main = expression(Omega ~ entries)
)
abline(0, 1, col = "red", lty = 2)

What success looks like. The fitted KL should be much smaller than at the start and close to the KL at par_true. Hold-out mean log-density under the fitted dp should nearly match the true SUN. Individual free parameters (especially \(L\) and the latent-correlation coordinates) may still differ when they are only weakly identified; agreement of the implied density and of Omega / Delta / Gamma is the more meaningful recovery check.

Parameter profiles: AD vs finite-difference gradients

Slice the cross-entropy objective \(-\sum_i w_i\log q(x_i)\) one coordinate at a time around the fitted value, and compare the AD partial derivative to a central finite-difference (Richardson) estimate on the same grid.

par0 <- as.numeric(fit$par)
names(par0) <- names(par_true)
w_fit <- as.numeric(fit$quad$w)
w_fit <- w_fit / sum(w_fit)

tape33 <- dsun_fun(
  x = as.matrix(fit$quad$x),
  par_seed = par0,
  weights = w_fit,
  n_points = 64L,
  n_shifts = 2L,
  n_threads = 1L
)

## dsun_fun value = sum w log q; objective = -sum w log q (= KL + const)
obj_ad <- function(p) {
  v <- tape33$eval(p, deriv = 0L)$value
  if (!is.finite(v)) {
    stop("non-finite AD objective at profile point", call. = FALSE)
  }
  -as.numeric(v)
}
grad_ad <- function(p) {
  g <- as.numeric(tape33$eval(p, deriv = 1L)$gradient)
  if (any(!is.finite(g))) {
    stop("non-finite AD gradient at profile point", call. = FALSE)
  }
  -g
}

bnd <- sun33_bounds()
col_kl <- "steelblue"

profile_seq <- function(j, n = 11L, half_width = 0.5) {
  center <- par0[j]
  lo <- bnd$lower[j]
  hi <- bnd$upper[j]
  hw <- half_width * max(abs(center), 1)
  lo_j <- center - hw
  hi_j <- center + hw
  if (is.finite(lo)) lo_j <- max(lo_j, lo)
  if (is.finite(hi)) hi_j <- min(hi_j, hi)
  if (!(is.finite(lo_j) && is.finite(hi_j) && lo_j < hi_j)) {
    stop(
      "Cannot build an interior profile sequence for parameter ", j,
      " (", names(par0)[j], ").",
      call. = FALSE
    )
  }
  seq(lo_j, hi_j, length.out = n)
}

## Plot obj + FD/AD gradient profiles for a block of parameter indices.
profile_par_block <- function(idx, n_col = 4L) {
  idx <- as.integer(idx)
  n_panels <- 2L * length(idx)
  n_row <- as.integer(ceiling(n_panels / n_col))
  op <- par(
    mfrow = c(n_row, n_col),
    mar = c(3, 3, 1.2, 0.4),
    mgp = c(1.6, 0.5, 0)
  )
  on.exit(par(op), add = TRUE)
  for (j in idx) {
    lab <- names(par0)[j]
    xj_seq <- profile_seq(j)
    obj_seq <- vapply(xj_seq, function(xj) {
      p <- par0
      p[j] <- xj
      obj_ad(p)
    }, numeric(1))
    ad_dobj <- vapply(xj_seq, function(xj) {
      p <- par0
      p[j] <- xj
      grad_ad(p)[j]
    }, numeric(1))
    ## Central FD at the same nodes as AD (Richardson on two step sizes).
    lo <- bnd$lower[j]
    hi <- bnd$upper[j]
    fd_at <- function(xj, h) {
      x_lo <- xj - h
      x_hi <- xj + h
      if (is.finite(lo)) x_lo <- max(x_lo, lo)
      if (is.finite(hi)) x_hi <- min(x_hi, hi)
      if (!(x_hi > x_lo)) {
        stop("central FD step collapsed for ", lab, " at ", format(xj))
      }
      p_lo <- p_hi <- par0
      p_lo[j] <- x_lo
      p_hi[j] <- x_hi
      (obj_ad(p_hi) - obj_ad(p_lo)) / (x_hi - x_lo)
    }
    fd_dobj <- vapply(xj_seq, function(xj) {
      h1 <- 1e-4 * max(abs(xj), 1)
      h2 <- 2 * h1
      fd1 <- fd_at(xj, h1)
      fd2 <- fd_at(xj, h2)
      (4 * fd1 - fd2) / 3
    }, numeric(1))

    obj_finite <- obj_seq[is.finite(obj_seq)]
    if (length(obj_finite)) {
      ylim_lo <- min(obj_finite)
      ylim_hi <- min(10, max(obj_finite))
      if (!(ylim_hi > ylim_lo)) {
        ylim_hi <- ylim_lo + 1e-8
      }
    } else {
      ylim_lo <- 0
      ylim_hi <- 10
    }
    plot(
      xj_seq, obj_seq,
      type = "b", pch = 16, col = col_kl, cex = 0.7,
      xlab = lab, ylab = expression(-sum(w * log(q))),
      ylim = c(ylim_lo, ylim_hi)
    )
    abline(v = par0[j], lty = 3, col = "grey60")
    plot(
      xj_seq, fd_dobj,
      type = "b", pch = 4, col = "grey30", cex = 1.4, lwd = 1.5,
      xlab = lab, ylab = paste0("d obj / d ", lab)
    )
    lines(xj_seq, ad_dobj, type = "b", pch = 3, col = col_kl, cex = 0.9, lwd = 1)
    abline(h = 0, lty = 3, col = "grey60")
    abline(v = par0[j], lty = 3, col = "grey60")
    legend(
      "topright",
      legend = c("FD", "AD"),
      col = c("grey30", col_kl), pch = c(4, 3),
      pt.cex = c(1.4, 0.9), bty = "n", cex = 0.65
    )
  }
}

profile_par_block(1:7)

SUN(3,3) slice profiles for parameters 1-7 (cross-entropy objective; FD vs AD gradient).

profile_par_block(8:14)

SUN(3,3) slice profiles for parameters 8-14.

profile_par_block(15:21)

SUN(3,3) slice profiles for parameters 15-21.

For each coordinate the left panel is the profiled objective and the right panel overlays the finite-difference slope (grey crosses) and the AD partial (blue pluses). Agreement of the two gradient curves is the gradient check; a vertical line marks the fitted value.