Embedded Laplace, solver 3: the Newton iteration converges to saddle points and converges slowly #3439

# issue_laplace_solver3.R
#
# Reproduction for a stan-dev/math issue on the Newton iteration of the
# embedded Laplace with solver 3 (LU), the solver that accepts a
# log-likelihood that is not concave.  Self-contained (the Stan programs are
# inline).
#
#   1. The Newton iteration converges to a saddle point of
#      log p(theta | y), where the Laplace approximation is not defined.
#   2. The iteration needs many steps: the line search starts from the
#      Barzilai-Borwein step and accepts steps up to about 1.9 times the
#      Newton step, so the convergence is linear, not quadratic.
#      Shown by the smallest max_num_steps that gives no
#      "max number of iterations" warning.
#
#   Rscript tests/issue_laplace_solver3.R
#   CMDSTAN_PATH=~/proj/cmdstan STAN_MATH=~/proj/stan_math_laplace_mean/ \
#     Rscript tests/issue_laplace_solver3.R

library(cmdstanr)
library(posterior)
if (nzchar(Sys.getenv("CMDSTAN_PATH"))) {
  set_cmdstan_path(Sys.getenv("CMDSTAN_PATH"))
}
cpp <- list()
if (nzchar(Sys.getenv("STAN_MATH"))) {
  cpp$MATH <- Sys.getenv("STAN_MATH")
}
compile <- function(code) {
  cmdstan_model(write_stan_file(code), cpp_options = cpp,
                force_recompile = length(cpp) > 0)
}
cat(sprintf("CmdStan %s at %s, Stan Math %s\n", cmdstan_version(),
            cmdstan_path(), if (is.null(cpp$MATH)) "of the CmdStan" else cpp$MATH))

# One fixed_param draw; returns the values of `vars` (NA if the draw
# failed) and the output lines.
run_once <- function(mod, data, vars) {
  out <- capture.output(
    fit <- mod$sample(data = data, fixed_param = TRUE, chains = 1,
                      iter_sampling = 1, refresh = 0, show_messages = TRUE,
                      sig_figs = 12),
    type = "output")
  vals <- tryCatch(as.numeric(as_draws_matrix(fit$draws(vars))),
                   error = function(e) rep(NA_real_, 2))
  list(vals = vals, out = out)
}
exception_of <- function(out) {
  unique(sub(".*Exception: (.*?) \\(in '.*", "\\1",
             grep("Exception", out, value = TRUE), perl = TRUE))
}

# ── 1. Saddle point ────────────────────────────────────────────────────────
saddle_code <- "
functions {
  real ll(vector theta, vector y, real s, real z, real sz) {
    return cauchy_lpdf(y | theta[1], s) + normal_lpdf(z | theta[2], sz);
  }
  matrix k_fun(matrix K) {
    return K;
  }
}
data {
  vector[2] y;
  real<lower=0> s;
  real z;
  real<lower=0> sz;
  matrix[2, 2] K;
  vector[2] theta0;
}
generated quantities {
  vector[2] mode;
  {
    tuple(vector[2], matrix[2, 2]) lap = laplace_latent_solve_tol(
      ll, (y, s, z, sz), 2, k_fun, (K,), (theta0, 1e-10, 100, 3, 100, 0));
    mode = lap.1;
  }
}
"
y <- c(-3, 3); s <- 0.3; z <- 0; sz <- 0.1
K <- matrix(c(4, 1, 1, 1), 2)
log_f1 <- function(th) {
  sum(dcauchy(y, th[1], s, log = TRUE)) + dnorm(z, th[2], sz, log = TRUE) -
    0.5 * sum(th * solve(K, th))
}
cat("\n1. Saddle point.  Stationary points of log p(theta | y):\n")
for (p in list(c(0, 0), c(2.9385, 0.0097), c(-2.9385, -0.0097))) {
  q <- optim(p, function(t) sum(numDeriv::grad(log_f1, t)^2), method = "BFGS",
             control = list(reltol = 1e-20))$par
  cat(sprintf("   (%7.4f, %7.4f): eigenvalues of -Hessian %s\n", q[1], q[2],
              paste(sprintf("%.3f", eigen(-numDeriv::hessian(log_f1, q),
                                          symmetric = TRUE)$values),
                    collapse = " ")))
}
mod1 <- compile(saddle_code)
for (start in list(c(0.3, 3), c(0.05, 3), c(0.01, 3), c(-0.01, 3),
                   c(-0.05, 3))) {
  r <- run_once(mod1, list(y = y, s = s, z = z, sz = sz, K = K,
                           theta0 = start), "mode")
  ex <- exception_of(r$out)
  cat(sprintf("   start (%5.2f, %4.1f): %s\n", start[1], start[2],
              if (length(ex)) paste("exception:", ex)
              else sprintf("mode (%.4f, %.4f)", r$vals[1], r$vals[2])))
}

# ── 2. Number of Newton iterations ─────────────────────────────────────────
tumor_code <- "
functions {
  real ll(vector theta, vector y, vector t, vector mu, real sigma) {
    real bsld = mu[1] * exp(theta[1]);
    real d = mu[2] * exp(theta[2]);
    real g = mu[3] * exp(theta[3]);
    real phi = inv_logit(logit(mu[4]) + theta[4]);
    vector[num_elements(y)] m
        = bsld * (phi * exp(-d * t) + (1 - phi) * exp(g * t));
    return normal_lpdf(y | m, sigma * m);
  }
  matrix k_fun(vector omega) {
    return diag_matrix(square(omega));
  }
}
data {
  int<lower=1> n;
  vector[n] y;
  vector[n] t;
  vector[4] mu;
  vector<lower=0>[4] omega;
  real<lower=0> sigma;
  vector[4] theta0;
  real<lower=0> tol;
  int<lower=1> max_steps;
}
generated quantities {
  vector[4] mode;
  {
    tuple(vector[4], matrix[4, 4]) lap = laplace_latent_solve_tol(
      ll, (y, t, mu, sigma), 4, k_fun, (omega,),
      (theta0, tol, max_steps, 3, 100, 0));
    mode = lap.1;
  }
}
"
data2 <- list(n = 12L,
              y = c(201.04, 78.30, 48.01, 44.39, 52.91, 50.49, 48.63, 67.29,
                    52.21, 72.43, 102.43, 105.06),
              t = seq(0, 693, by = 63),
              mu = c(117.4, 9.121e-3, 1.913e-3, 0.3781),
              omega = c(0.4925, 0.9217, 1.120, 1.299), sigma = 0.09826,
              theta0 = c(0.3636, -0.3385, 1.9794, -1.9322))
mod2 <- compile(tumor_code)
hits_limit <- function(tol, max_steps) {
  r <- run_once(mod2, c(data2, list(tol = tol, max_steps = as.integer(max_steps))),
                "mode")
  any(grepl("max number of iterations", r$out))
}
cat("\n2. Newton iterations, 4 latent variables (smallest max_num_steps without the\n",
    "   'max number of iterations' warning):\n")
for (tol in c(1e-6, 1.49012e-8, 1e-10)) {
  lo <- 1L
  hi <- 200L
  while (lo < hi) {
    mid <- (lo + hi) %/% 2L
    if (hits_limit(tol, mid)) lo <- mid + 1L else hi <- mid
  }
  cat(sprintf("   tol %-11g max_num_steps %d\n", tol, lo))
}