cleanup tests and address review comments

Author: SteveBronder

stan/math/mix/functor/laplace_marginal_density_estimator.hpp

@@ -444,7 +444,7 @@ inline void llt_with_jitter(LLT& llt_B, B_t& B, double min_jitter = 1e-10,
     }
     if (llt_B.info() != Eigen::Success) {
       throw std::domain_error(
-          "laplace_marginal_density: Cholesky (Diag) failed");
+          "laplace_marginal_density: Cholesky failed after adding jitter up to " + std::to_string(jitter_try));
     }
   }
 }
@@ -942,16 +942,12 @@ inline auto run_newton_loop(SolverPolicy& solver, NewtonStateT& state,
       scratch.alpha() = 1.0;
       update_fun(scratch, state.curr(), state.prev(), scratch.eval_,
                  state.wolfe_info.p_);
-      bool run_convergence_check = true;
       if (scratch.alpha() <= options.line_search.min_alpha) {
         state.wolfe_status.accept_ = false;
         finish_update = true;
-        run_convergence_check = false;
       } else if (options.line_search.max_iterations == 0) {
         state.curr().update(scratch);
         state.wolfe_status.accept_ = true;
-        finish_update = false;
-        run_convergence_check = false;
       } else {
         Eigen::VectorXd s = scratch.a() - state.prev().a();
         auto full_step_grad
@@ -964,16 +960,15 @@ inline auto run_newton_loop(SolverPolicy& solver, NewtonStateT& state,
         state.wolfe_status = internal::wolfe_line_search(
             state.wolfe_info, update_fun, options.line_search, msgs);
       }
-      if (run_convergence_check) {
-        /**
-         * Stop when objective change is small, or when a rejected Wolfe step
-         * fails to improve; finish_update then exits the Newton loop.
-         */
-        finish_update = std::abs(state.curr().obj() - state.prev().obj())
-                            < options.tolerance
-                        || (!state.wolfe_status.accept_
-                            && state.curr().obj() <= state.prev().obj());
-      }
+      /**
+       * Stop when objective change is small, or when a rejected Wolfe step
+       * fails to improve; finish_update then exits the Newton loop.
+       */
+      bool objective_converged = std::abs(state.curr().obj() - state.prev().obj())
+                          < options.tolerance;
+      bool search_failed = (!state.wolfe_status.accept_
+                          && state.curr().obj() <= state.prev().obj());
+      finish_update = objective_converged || search_failed;
     }
     if (finish_update) {
       if (!state.final_loop && state.wolfe_status.accept_) {
@@ -1039,6 +1034,52 @@ inline decltype(auto) theta_init_impl(Eigen::Index theta_size, Opts&& options) {
   }
 }
 
+/**
+ * @brief Create the update function for the line search, capturing necessary
+ * references.
+ * @tparam ObjFun Callable type for the objective function (accepting (a, theta))
+ * @tparam ThetaGradFun Callable type for the theta gradient function (accepting theta)
+ * @tparam Covariance Type of the covariance matrix
+ * @tparam Options Type of the options struct containing line search parameters
+ * @param[in] obj_fun Objective function functor
+ * @param[in] theta_grad_f Theta gradient functor
+ * @param[in] covariance Prior covariance matrix Sigma
+ * @param[in] options Options struct containing line search parameters
+ * @return A callable update function for the line search, with signature:
+ * ```
+ *  bool update_fun(proposal, curr, prev, eval_in, p)
+ * ```
+ */
+template <typename ObjFun, typename ThetaGradFun, typename Covariance, typename Options>
+inline auto create_update_fun(ObjFun&& obj_fun, ThetaGradFun&& theta_grad_f,
+                             Covariance&& covariance, Options&& options) {
+  auto update_step = [&covariance, &obj_fun, &theta_grad_f](
+                          auto& proposal, auto&& /* curr */, auto&& prev,
+                          auto& eval_in, auto&& p) {
+    try {
+      proposal.a() = prev.a() + eval_in.alpha() * p;
+      proposal.theta().noalias() = covariance * proposal.a();
+      proposal.theta_grad() = theta_grad_f(proposal.theta());
+      eval_in.obj() = obj_fun(proposal.a(), proposal.theta());
+      eval_in.dir()
+          = (-covariance * proposal.a() + covariance * proposal.theta_grad())
+                .dot(p);
+      return std::isfinite(eval_in.obj()) && std::isfinite(eval_in.dir());
+    } catch (const std::exception&) {
+      return false;
+    }
+  };
+  auto backoff = [&options](auto& eval) {
+    eval.alpha() *= options.line_search.tau;
+    return eval.alpha() > options.line_search.min_alpha;
+  };
+  return [update_step_ = std::move(update_step), backoff_ = std::move(backoff)](
+          auto& proposal, auto&& curr, auto&& prev, auto& eval_in, auto&& p) {
+        return internal::retry_evaluate(update_step_, proposal, curr, prev,
+                                        eval_in, p, backoff_);
+      };
+}
+
 /**
  * For a latent Gaussian model with hyperparameters phi and
  * latent variables theta, and observations y, this function computes
@@ -1098,9 +1139,7 @@ inline auto laplace_marginal_density_est(
     const laplace_options<InitTheta>& options, std::ostream* msgs) {
   internal::validate_laplace_options("laplace_marginal_density", options,
                                      covariance);
-
   const Eigen::Index theta_size = covariance.rows();
-
   // Wolfe optimizes over the latent 'a' space
   auto obj_fun = [&ll_fun, &ll_args, &msgs](const Eigen::VectorXd& a_val,
                                             auto&& theta_val) -> double {
@@ -1113,36 +1152,8 @@ inline auto laplace_marginal_density_est(
   };
   decltype(auto) theta_init = theta_init_impl<InitTheta>(theta_size, options);
   internal::NewtonState state(theta_size, obj_fun, theta_grad_f, theta_init);
-  // 'a' gradient
-  auto update_fun = [&covariance, &obj_fun, &theta_grad_f, &options]() {
-    auto update_step = [&covariance, &obj_fun, &theta_grad_f](
-                           auto& proposal, auto&& /* curr */, auto&& prev,
-                           auto& eval_in, auto&& p) {
-      try {
-        proposal.a() = prev.a() + eval_in.alpha() * p;
-        proposal.theta().noalias() = covariance * proposal.a();
-        proposal.theta_grad() = theta_grad_f(proposal.theta());
-        eval_in.obj() = obj_fun(proposal.a(), proposal.theta());
-        eval_in.dir()
-            = (-covariance * proposal.a() + covariance * proposal.theta_grad())
-                  .dot(p);
-        return std::isfinite(eval_in.obj()) && std::isfinite(eval_in.dir());
-      } catch (const std::exception&) {
-        return false;
-      }
-    };
-    auto backoff = [&options](auto& eval) {
-      eval.alpha() *= options.line_search.tau;
-      return eval.alpha() > options.line_search.min_alpha;
-    };
-    return
-        [update_step_ = std::move(update_step), backoff_ = std::move(backoff)](
-            auto& proposal, auto&& curr, auto&& prev, auto& eval_in, auto&& p) {
-          return internal::retry_evaluate(update_step_, proposal, curr, prev,
-                                          eval_in, p, backoff_);
-        };
-  }();
   // Start with safe step size
+  auto update_fun = create_update_fun(std::move(obj_fun), std::move(theta_grad_f), covariance, options);
   Eigen::Index step_iter = 0;
   try {
     if (options.solver == 1) {