laplace_marginal_density_estimator.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_MIX_FUNCTOR_LAPLACE_MARGINAL_DENSITY_ESTIMATOR_HPP
#define STAN_MATH_MIX_FUNCTOR_LAPLACE_MARGINAL_DENSITY_ESTIMATOR_HPP
#include <stan/math/prim/fun/Eigen.hpp>
#include <stan/math/prim/fun/generate_laplace_options.hpp>
#include <stan/math/mix/functor/laplace_likelihood.hpp>
#include <stan/math/mix/functor/wolfe_line_search.hpp>
#include <stan/math/mix/functor/barzilai_borwein_step_size.hpp>
#include <stan/math/mix/meta.hpp>
#include <stan/math/prim/fun/to_ref.hpp>
#include <stan/math/prim/fun/quad_form_diag.hpp>
#include <stan/math/prim/fun/value_of.hpp>
#include <stan/math/prim/functor/iter_tuple_nested.hpp>
#include <algorithm>
#include <cmath>
#include <limits>
#include <mutex>
#include <iomanip>
 
namespace stan {
namespace math {
 
struct laplace_options_base {
  /* Size of the blocks in block diagonal hessian*/
  int hessian_block_size{internal::laplace_default_hessian_block_size};  // 0
  int solver{internal::laplace_default_solver};  // 1
  double tolerance{internal::laplace_default_tolerance};  // 2
  /* Maximum number of steps*/
  int max_num_steps{internal::laplace_default_max_num_steps};          // 3
  int allow_fallthrough{internal::laplace_default_allow_fallthrough};  // 4
  laplace_line_search_options line_search{
      internal::laplace_default_max_steps_line_search};  // 5
  laplace_options_base() = default;
  laplace_options_base(int hessian_block_size_, int solver_, double tolerance_,
                       int max_num_steps_, bool allow_fallthrough_,
                       int max_steps_line_search_)
      : hessian_block_size(hessian_block_size_),
        solver(solver_),
        tolerance(tolerance_),
        max_num_steps(max_num_steps_),
        allow_fallthrough(allow_fallthrough_),
        line_search(max_steps_line_search_) {}
};
 
template <bool HasInitTheta>
struct laplace_options;
 
template <>
struct laplace_options<false> : public laplace_options_base {
  laplace_options() = default;
 
  explicit laplace_options(int hessian_block_size_) {
    hessian_block_size = hessian_block_size_;
  }
};
 
template <>
struct laplace_options<true> : public laplace_options_base {
  /* Value for user supplied initial theta  */
  Eigen::VectorXd theta_0{0};  // 6
 
  template <typename ThetaVec>
  laplace_options(ThetaVec&& theta_0_, double tolerance_, int max_num_steps_,
                  int hessian_block_size_, int solver_,
                  int max_steps_line_search_, bool allow_fallthrough_)
      : laplace_options_base(hessian_block_size_, solver_, tolerance_,
                             max_num_steps_, allow_fallthrough_,
                             max_steps_line_search_),
        theta_0(value_of(std::forward<ThetaVec>(theta_0_))) {}
};
 
using laplace_options_default = laplace_options<false>;
using laplace_options_user_supplied = laplace_options<true>;
 
namespace internal {
 
template <typename Options>
inline constexpr auto tuple_to_laplace_options(Options&& ops) {
  using Ops = std::decay_t<Options>;
  if constexpr (is_tuple_v<Ops>) {
    if constexpr (!is_eigen_v<std::tuple_element_t<0, std::decay_t<Ops>>>) {
      static_assert(
          sizeof(std::decay_t<Ops>*) == 0,
          "ERROR:(laplace_marginal_lpdf) The first laplace argument is "
          "expected to be an Eigen vector of dynamic size representing the "
          "initial theta_0.");
    }
    if constexpr (!stan::is_inner_tuple_type_v<1, Ops, double>) {
      static_assert(
          sizeof(std::decay_t<Ops>*) == 0,
          "ERROR:(laplace_marginal_lpdf) The second laplace argument is "
          "expected to be a double representing the tolerance.");
    }
    if constexpr (!stan::is_inner_tuple_type_v<2, Ops, int>) {
      static_assert(
          sizeof(std::decay_t<Ops>*) == 0,
          "ERROR:(laplace_marginal_lpdf) The third laplace argument is "
          "expected to be an int representing the maximum number of steps for "
          "the laplace approximation.");
    }
    if constexpr (!stan::is_inner_tuple_type_v<3, Ops, int>) {
      static_assert(
          sizeof(std::decay_t<Ops>*) == 0,
          "ERROR:(laplace_marginal_lpdf) The fourth laplace argument is "
          "expected to be an int representing the solver.");
    }
    if constexpr (!stan::is_inner_tuple_type_v<4, Ops, int>) {
      static_assert(
          sizeof(std::decay_t<Ops>*) == 0,
          "ERROR:(laplace_marginal_lpdf) The fifth laplace argument is "
          "expected to be an int representing the max steps for the laplace "
          "approximaton's wolfe line search.");
    }
    constexpr bool is_fallthrough
        = stan::is_inner_tuple_type_v<
              5, Ops, int> || stan::is_inner_tuple_type_v<5, Ops, bool>;
    if constexpr (!is_fallthrough) {
      static_assert(
          sizeof(std::decay_t<Ops>*) == 0,
          "ERROR:(laplace_marginal_lpdf) The sixth laplace argument is "
          "expected to be an int representing allow fallthrough (0/1).");
    }
    auto defaults = laplace_options_default{};
    return laplace_options_user_supplied{
        value_of(std::get<0>(std::forward<Options>(ops))),
        std::get<1>(ops),
        std::get<2>(ops),
        defaults.hessian_block_size,
        std::get<3>(ops),
        std::get<4>(ops),
        (std::get<5>(ops) > 0) ? true : false,
    };
  } else {
    return std::forward<Options>(ops);
  }
}
 
template <typename ThetaVec, typename WR, typename L_t, typename A_vec,
          typename ThetaGrad, typename LU_t, typename KRoot>
struct laplace_density_estimates {
  /* log marginal density */
  double lmd{std::numeric_limits<double>::infinity()};
  /* ThetaVec at the mode */
  ThetaVec theta;
  WR W_r;
  L_t L;
  A_vec a;
  ThetaGrad theta_grad;
  /* LU matrix from solver 3 */
  LU_t LU;
  KRoot K_root;
  int solver_used{1};
  laplace_density_estimates(double lmd_, ThetaVec&& theta_, WR&& W_r_, L_t&& L_,
                            A_vec&& a_, ThetaGrad&& theta_grad_, LU_t&& LU_,
                            KRoot&& K_root_, int solver_used_)
      : lmd(lmd_),
        theta(std::move(theta_)),
        W_r(std::move(W_r_)),
        L(std::move(L_)),
        a(std::move(a_)),
        theta_grad(std::move(theta_grad_)),
        LU(std::move(LU_)),
        K_root(std::move(K_root_)),
        solver_used(solver_used_) {}
};
 
template <typename WRootMat>
inline void block_matrix_sqrt(WRootMat& W_root,
                              const Eigen::SparseMatrix<double>& W,
                              const Eigen::Index block_size) {
  const Eigen::Index n_block = W.cols() / block_size;
  Eigen::MatrixXd local_block(block_size, block_size);
  Eigen::MatrixXd local_block_sqrt(block_size, block_size);
  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eigensolver;
  // No block operation available for sparse matrices, so we have to loop
  // See https://eigen.tuxfamily.org/dox/group__TutorialSparse.html#title7
  for (Eigen::Index i = 0; i < n_block; i++) {
    local_block
        = W.block(i * block_size, i * block_size, block_size, block_size);
    if (unlikely(!local_block.array().isFinite().all())) {
      [](auto i) STAN_COLD_PATH {
        throw std::domain_error(
            std::string("Error in block_matrix_sqrt: "
                        "non-finite values detected in block diagonal "
                        "starting at (")
            + std::to_string(i) + ", " + std::to_string(i) + ")");
      }(i);
    }
    local_block_sqrt = 0.5 * (local_block + local_block.transpose());
    eigensolver.compute(local_block_sqrt);
    if (unlikely(eigensolver.info() != Eigen::Success)) {
      [](auto i) STAN_COLD_PATH {
        throw std::domain_error(
            std::string("Error in block_matrix_sqrt: "
                        "eigendecomposition failed for block diagonal "
                        "starting at (")
            + std::to_string(i) + ", " + std::to_string(i) + ")");
      }(i);
    }
    const Eigen::VectorXd eigenvalues = eigensolver.eigenvalues();
    const double tolerance = block_size * std::numeric_limits<double>::epsilon()
                             * eigenvalues.cwiseAbs().maxCoeff();
    if (unlikely(eigenvalues.minCoeff() < -tolerance)) {
      [](auto&& i, auto&& eigenvalues) {
        throw std::domain_error(
            std::string("Error in block_matrix_sqrt: block diagonal starting "
                        "at (")
            + std::to_string(i) + ", " + std::to_string(i)
            + ") is not positive semi-definite (smallest eigenvalue "
            + std::to_string(eigenvalues.minCoeff()) + ")");
      }(i, eigenvalues);
    }
    local_block_sqrt.noalias()
        = eigensolver.eigenvectors()
          * eigenvalues.cwiseMax(0.0).cwiseSqrt().asDiagonal()
          * eigensolver.eigenvectors().transpose();
    for (Eigen::Index k = 0; k < block_size; k++) {
      for (Eigen::Index j = 0; j < block_size; j++) {
        W_root.coeffRef(i * block_size + j, i * block_size + k)
            = local_block_sqrt(j, k);
      }
    }
  }
}
 
template <bool InitTheta, typename CovarMat>
inline void validate_laplace_options(const char* frame_name,
                                     const laplace_options<InitTheta>& options,
                                     const CovarMat& covariance) {
  if constexpr (InitTheta) {
    check_nonzero_size(frame_name, "initial guess", options.theta_0);
    check_finite(frame_name, "initial guess", options.theta_0);
    if (unlikely(options.theta_0.size() != covariance.rows())) {
      std::stringstream msg;
      msg << frame_name << ": The size of the initial theta ("
          << options.theta_0.size()
          << ") vector must match the rows and columns of the covariance "
             "matrix ("
          << covariance.rows() << ", " << covariance.cols() << ").";
      throw std::domain_error(msg.str());
    }
  }
  check_nonnegative(frame_name, "tolerance", options.tolerance);
  check_positive(frame_name, "max_num_steps", options.max_num_steps);
  check_positive(frame_name, "hessian_block_size", options.hessian_block_size);
  check_square(frame_name, "covariance", covariance);
 
  const Eigen::Index theta_size = covariance.rows();
  if (unlikely(theta_size % options.hessian_block_size != 0
               || theta_size < options.hessian_block_size)) {
    throw std::domain_error(
        "laplace_marginal_density: Hessian block size mismatch.");
  }
 
  if (unlikely(options.solver < 1 || options.solver > 3)) {
    throw std::domain_error(
        "laplace_marginal_density: solver must be 1, 2, or 3. Got: "
        + std::to_string(options.solver));
  }
}
 
struct NewtonState {
  WolfeInfo wolfe_info;
 
  WolfeStatus wolfe_status;
 
  WolfeData proposal;
 
  Eigen::VectorXd b;
 
  Eigen::MatrixXd B;
 
  Eigen::VectorXd prev_g;
  bool final_loop = false;
 
  template <typename ObjFun, typename ThetaGradFun, typename CovarianceT,
            typename ThetaInitializer>
  NewtonState(int theta_size, ObjFun&& obj_fun, ThetaGradFun&& theta_grad_f,
              CovarianceT&& covariance, ThetaInitializer&& theta_init)
      : wolfe_info(std::forward<ObjFun>(obj_fun),
                   covariance.llt().solve(theta_init),
                   std::forward<ThetaInitializer>(theta_init),
                   std::forward<ThetaGradFun>(theta_grad_f)),
        proposal(theta_size),
        b(theta_size),
        B(theta_size, theta_size),
        prev_g(theta_size) {
    wolfe_status.num_backtracks_ = -1;  // Safe initial value for BB step
  }
 
  auto& curr() & { return wolfe_info.curr_; }
 
  const auto& curr() const& { return wolfe_info.curr_; }
  auto&& curr() && { return std::move(wolfe_info).curr(); }
  auto& prev() & { return wolfe_info.prev_; }
 
  const auto& prev() const& { return wolfe_info.prev_; }
  auto&& prev() && { return std::move(wolfe_info).prev(); }
  auto& proposal_step() & { return proposal; }
  const auto& proposal_step() const& { return proposal; }
  auto&& proposal_step() && { return std::move(proposal); }
  template <typename Options>
  inline void update_next_step(const Options& options) {
    this->prev().swap(this->curr());
    this->curr().alpha()
        = std::clamp(this->curr().alpha(), 0.0, options.line_search.max_alpha);
  }
};
 
template <typename LLT, typename B_t>
inline void llt_with_jitter(LLT& llt_B, B_t& B, double min_jitter = 1e-10,
                            double max_jitter = 1e-5) {
  llt_B.compute(B);
  if (llt_B.info() != Eigen::Success) {
    double prev_jitter = 0.0;
    double jitter_try = min_jitter;
    for (; jitter_try < max_jitter; jitter_try *= 10) {
      // Remove previously added jitter before adding the new (larger) amount,
      // so that the total diagonal perturbation is exactly jitter_try.
      B.diagonal().array() += (jitter_try - prev_jitter);
      prev_jitter = jitter_try;
      llt_B.compute(B);
      if (llt_B.info() == Eigen::Success) {
        break;
      }
    }
    if (llt_B.info() != Eigen::Success) {
      throw std::domain_error(
          "laplace_marginal_density: Cholesky failed after adding jitter up to "
          + std::to_string(jitter_try));
    }
  }
}
 
struct CholeskyWSolverDiag {
  Eigen::VectorXd W_r_diag;
 
  Eigen::VectorXd W_diag;
 
  Eigen::LLT<Eigen::MatrixXd> llt_B;
 
  template <typename NewtonStateT, typename CovarMat>
  CholeskyWSolverDiag(const NewtonStateT& state, const CovarMat& covariance)
      : W_r_diag(Eigen::VectorXd::Zero(state.b.size())), W_diag(0), llt_B() {}
  template <typename NewtonStateT, typename LLFun, typename LLTupleArgs,
            typename CovarMat>
  void solve_step(NewtonStateT& state, const LLFun& ll_fun,
                  const LLTupleArgs& ll_args, const CovarMat& covariance,
                  int /*hessian_block_size*/, std::ostream* msgs) {
    const Eigen::Index theta_size = state.b.size();
 
    // 1. Compute diagonal Hessian
    W_diag = laplace_likelihood::diagonal_hessian(ll_fun, state.prev().theta(),
                                                  ll_args, msgs);
    for (Eigen::Index j = 0; j < W_diag.size(); j++) {
      if (W_diag.coeff(j) < 0 || !std::isfinite(W_diag.coeff(j))) {
        throw std::domain_error(
            "laplace_marginal_density: Hessian matrix is not positive "
            "definite");
      } else {
        W_r_diag.coeffRef(j) = std::sqrt(W_diag.coeff(j));
      }
    }
 
    // 2. Formulate B = I + W_r * Sigma * W_r
    state.B.noalias()
        = Eigen::MatrixXd::Identity(theta_size, theta_size)
          + W_r_diag.asDiagonal() * covariance * W_r_diag.asDiagonal();
 
    // 3. Factorize B with jittering fallback
    llt_with_jitter(llt_B, state.B);
    // 4. Solve for the raw Newton proposal in a-space.
    state.b.noalias() = (W_diag.array() * state.prev().theta().array()).matrix()
                        + state.prev().theta_grad();
    auto L = llt_B.matrixL();
    auto LT = llt_B.matrixU();
    state.proposal_step().a().noalias()
        = state.b
          - W_r_diag.asDiagonal()
                * LT.solve(
                    L.solve(W_r_diag.cwiseProduct(covariance * state.b)));
  }
 
  double compute_log_determinant() const {
    return 2.0 * llt_B.matrixLLT().diagonal().array().log().sum();
  }
 
  template <typename NewtonStateT>
  auto build_result(NewtonStateT& state, double log_det) {
    return laplace_density_estimates{
        state.prev().obj() - 0.5 * log_det,
        std::move(state).prev().theta(),
        Eigen::SparseMatrix<double>(W_r_diag.asDiagonal()),
        Eigen::MatrixXd(llt_B.matrixL()),
        std::move(state).prev().a(),
        std::move(state).prev().theta_grad(),
        Eigen::PartialPivLU<Eigen::MatrixXd>{},
        Eigen::MatrixXd(0, 0),
        1};
  }
};
 
struct CholeskyWSolverBlock {
  Eigen::SparseMatrix<double> W_r;
 
  Eigen::SparseMatrix<double> W_block;
 
  Eigen::LLT<Eigen::MatrixXd> llt_B;
 
  template <typename NewtonStateT>
  CholeskyWSolverBlock(const NewtonStateT& state, int hessian_block_size)
      : W_r(state.b.size(), state.b.size()) {
    const Eigen::Index theta_size = state.b.size();
    W_r.reserve(Eigen::VectorXi::Constant(theta_size, hessian_block_size));
    const Eigen::Index n_block = theta_size / hessian_block_size;
    for (Eigen::Index ii = 0; ii < n_block; ii++) {
      for (Eigen::Index k = 0; k < hessian_block_size; k++) {
        for (Eigen::Index j = 0; j < hessian_block_size; j++) {
          W_r.insert(ii * hessian_block_size + j, ii * hessian_block_size + k)
              = 1.0;
        }
      }
    }
    W_r.makeCompressed();
  }
 
  template <typename NewtonStateT, typename LLFun, typename LLTupleArgs,
            typename CovarMat>
  void solve_step(NewtonStateT& state, const LLFun& ll_fun,
                  const LLTupleArgs& ll_args, const CovarMat& covariance,
                  int hessian_block_size, std::ostream* msgs) {
    const Eigen::Index theta_size = state.b.size();
    // 1. Compute block Hessian
    W_block = laplace_likelihood::block_hessian(
        ll_fun, state.prev().theta(), hessian_block_size, ll_args, msgs);
 
    for (Eigen::Index j = 0; j < W_block.rows(); j++) {
      if (W_block.coeff(j, j) < 0 || !std::isfinite(W_block.coeff(j, j))) {
        throw std::domain_error(
            "laplace_marginal_density: Hessian matrix is not positive "
            "definite");
      }
    }
 
    // 2. Compute W_r = sqrt(W)
    block_matrix_sqrt(W_r, W_block, hessian_block_size);
 
    // 3. Formulate B = I + W_r * Sigma * W_r
    state.B.noalias() = Eigen::MatrixXd::Identity(theta_size, theta_size)
                        + W_r * (covariance * W_r);
 
    // 4. Factorize B with jittering fallback
    llt_with_jitter(llt_B, state.B);
 
    // 5. Solve for the raw Newton proposal in a-space.
    state.b.noalias()
        = W_block * state.prev().theta() + state.prev().theta_grad();
    auto L = llt_B.matrixL();
    auto LT = llt_B.matrixU();
    state.proposal_step().a().noalias()
        = state.b - W_r * LT.solve(L.solve(W_r * (covariance * state.b)));
  }
 
  double compute_log_determinant() const {
    return 2.0 * llt_B.matrixLLT().diagonal().array().log().sum();
  }
 
  template <typename NewtonStateT>
  auto build_result(NewtonStateT& state, double log_det) {
    return laplace_density_estimates{state.prev().obj() - 0.5 * log_det,
                                     std::move(state).prev().theta(),
                                     std::move(W_r),
                                     Eigen::MatrixXd(llt_B.matrixL()),
                                     std::move(state).prev().a(),
                                     std::move(state).prev().theta_grad(),
                                     Eigen::PartialPivLU<Eigen::MatrixXd>{},
                                     Eigen::MatrixXd(0, 0),
                                     1};
  }
};
 
struct CholeskyKSolver {
  Eigen::MatrixXd K_root;
 
  Eigen::SparseMatrix<double> W_full;
 
  Eigen::LLT<Eigen::MatrixXd> llt_B;
 
  template <typename NewtonStateT, typename CovarMat>
  CholeskyKSolver(const NewtonStateT& state, const CovarMat& covariance)
      : K_root(0, 0), W_full(0, 0), llt_B() {
    auto K_root_llt = covariance.template selfadjointView<Eigen::Lower>().llt();
    if (K_root_llt.info() != Eigen::Success) {
      throw std::domain_error(
          "laplace_marginal_density: Cholesky of covariance failed at start");
    }
    K_root = std::move(K_root_llt.matrixL());
  }
 
  template <typename NewtonStateT, typename LLFun, typename LLTupleArgs,
            typename CovarMat>
  void solve_step(NewtonStateT& state, const LLFun& ll_fun,
                  const LLTupleArgs& ll_args, const CovarMat& covariance,
                  int hessian_block_size, std::ostream* msgs) {
    const Eigen::Index theta_size = state.b.size();
 
    // 1. Compute Hessian
    W_full = laplace_likelihood::block_hessian(
        ll_fun, state.prev().theta(), hessian_block_size, ll_args, msgs);
 
    // 2. Formulate B = I + K^T * W * K
    state.B.noalias() = Eigen::MatrixXd::Identity(theta_size, theta_size)
                        + K_root.transpose() * (W_full * K_root);
 
    // 3. Factorize B with jittering fallback
    llt_with_jitter(llt_B, state.B);
 
    // 4. Solve for the raw Newton proposal in a-space.
    state.b.noalias()
        = W_full * state.prev().theta() + state.prev().theta_grad();
    auto L = llt_B.matrixL();
    auto LT = llt_B.matrixU();
    state.proposal_step().a().noalias()
        = K_root.transpose().template triangularView<Eigen::Upper>().solve(
            LT.solve(L.solve(K_root.transpose() * state.b)));
  }
 
  double compute_log_determinant() const {
    return 2.0 * llt_B.matrixLLT().diagonal().array().log().sum();
  }
 
  template <typename NewtonStateT>
  auto build_result(NewtonStateT& state, double log_det) {
    return laplace_density_estimates{state.prev().obj() - 0.5 * log_det,
                                     std::move(state.prev().theta()),
                                     std::move(W_full),
                                     Eigen::MatrixXd(llt_B.matrixL()),
                                     std::move(state.prev().a()),
                                     std::move(state.prev().theta_grad()),
                                     Eigen::PartialPivLU<Eigen::MatrixXd>{},
                                     std::move(K_root),
                                     2};
  }
};
 
struct LUSolver {
  Eigen::PartialPivLU<Eigen::MatrixXd> lu;
 
  Eigen::SparseMatrix<double> W_full;
 
  template <typename NewtonStateT, typename LLFun, typename LLTupleArgs,
            typename CovarMat>
  void solve_step(NewtonStateT& state, const LLFun& ll_fun,
                  const LLTupleArgs& ll_args, const CovarMat& covariance,
                  int hessian_block_size, std::ostream* msgs) {
    const Eigen::Index theta_size = state.b.size();
 
    // 1. Compute Hessian
    W_full = laplace_likelihood::block_hessian(
        ll_fun, state.prev().theta(), hessian_block_size, ll_args, msgs);
 
    // 2. Factorize B = I + Sigma * W
    lu.compute(Eigen::MatrixXd::Identity(theta_size, theta_size)
               + covariance * W_full);
 
    // 3. Solve for the raw Newton proposal in a-space.
    state.b.noalias()
        = W_full * state.prev().theta() + state.prev().theta_grad();
    state.proposal_step().a().noalias()
        = state.b - W_full * lu.solve(covariance * state.b);
  }
 
  double compute_log_determinant() const {
    return lu.matrixLU().diagonal().array().log().sum();
  }
 
  template <typename NewtonStateT>
  auto build_result(NewtonStateT& state, double log_det) {
    return laplace_density_estimates{state.prev().obj() - 0.5 * log_det,
                                     std::move(state).prev().theta(),
                                     std::move(W_full),
                                     Eigen::MatrixXd(0, 0),
                                     std::move(state).prev().a(),
                                     std::move(state).prev().theta_grad(),
                                     std::move(lu),
                                     Eigen::MatrixXd(0, 0),
                                     3};
  }
};
 
template <typename SolverPolicy, typename NewtonStateT, typename OptionsT,
          typename LLFunT, typename LLTupleArgsT, typename CovarMatT,
          typename UpdateFun>
inline auto run_newton_loop(SolverPolicy& solver, NewtonStateT& state,
                            const OptionsT& options, Eigen::Index& step_iter,
                            const LLFunT& ll_fun, const LLTupleArgsT& ll_args,
                            const CovarMatT& covariance, UpdateFun&& update_fun,
                            std::ostream* msgs) {
  bool finish_update = false;
  for (; step_iter <= options.max_num_steps; step_iter++) {
    solver.solve_step(state, ll_fun, ll_args, covariance,
                      options.hessian_block_size, msgs);
    if (!state.final_loop) {
      auto&& proposal = state.proposal_step();
      state.wolfe_info.p_ = proposal.a() - state.prev().a();
      state.prev_g.noalias() = -covariance * state.prev().a()
                               + covariance * state.prev().theta_grad();
      state.wolfe_info.init_dir_ = state.prev_g.dot(state.wolfe_info.p_);
      // Flip direction if not ascending
      state.wolfe_info.flip_direction();
      auto&& scratch = state.wolfe_info.scratch_;
      proposal.eval_.alpha() = 1.0;
      const bool proposal_valid
          = update_fun(proposal, state.curr(), state.prev(), proposal.eval_,
                       state.wolfe_info.p_);
      const bool cached_proposal_ok
          = proposal_valid && std::isfinite(proposal.obj())
            && std::isfinite(proposal.dir())
            && proposal.alpha() > options.line_search.min_alpha;
      if (!cached_proposal_ok) {
        state.wolfe_status
            = WolfeStatus{WolfeReturn::StepTooSmall, 1, 0, false};
      } else if (options.line_search.max_iterations == 0) {
        state.curr().update(proposal);
        state.wolfe_status = WolfeStatus{WolfeReturn::Continue, 1, 0, true};
      } else {
        Eigen::VectorXd s = proposal.a() - state.prev().a();
        auto full_step_grad
            = (-covariance * proposal.a() + covariance * proposal.theta_grad())
                  .eval();
        state.curr().alpha() = barzilai_borwein_step_size(
            s, full_step_grad, state.prev_g, state.prev().alpha(),
            state.wolfe_status.num_backtracks_, options.line_search.min_alpha,
            options.line_search.max_alpha);
        state.wolfe_status = internal::wolfe_line_search(
            state.wolfe_info, update_fun, options.line_search, msgs);
      }
      bool search_failed = !state.wolfe_status.accept_;
      const bool proposal_armijo_ok
          = cached_proposal_ok
            && internal::check_armijo(
                proposal.obj(), state.prev().obj(), proposal.alpha(),
                state.wolfe_info.init_dir_, options.line_search);
      if (search_failed && proposal_armijo_ok) {
        state.curr().update(proposal);
        state.wolfe_status
            = WolfeStatus{WolfeReturn::Armijo, state.wolfe_status.num_evals_,
                          state.wolfe_status.num_backtracks_, true};
        search_failed = false;
      }
      bool objective_converged
          = state.wolfe_status.accept_
            && std::abs(state.curr().obj() - state.prev().obj())
                   < options.tolerance;
      finish_update = objective_converged || search_failed;
    }
    if (finish_update) {
      if (!state.final_loop && state.wolfe_status.accept_) {
        // Do one final loop with exact wolfe conditions
        state.final_loop = true;
        state.update_next_step(options);
        continue;
      }
      return solver.build_result(state, solver.compute_log_determinant());
    } else {
      state.update_next_step(options);
    }
  }
  if (msgs) {
    (*msgs)
        << std::string(
               "WARNING(laplace_marginal_density): max number of iterations: ")
               + std::to_string(options.max_num_steps) + " exceeded.";
  }
  return solver.build_result(state, solver.compute_log_determinant());
}
 
[[noreturn]] inline void throw_solver_failure(std::string_view context,
                                              Eigen::Index iter,
                                              std::string_view failed_solver,
                                              const std::exception& e) {
  std::ostringstream os;
  os << context << ": " << failed_solver << " failed at iteration " << iter
     << " and allow_fallthrough is false. Reason: " << e.what();
  throw std::domain_error(os.str());
}
 
inline void log_solver_fallback(std::ostream* msgs, std::string_view context,
                                Eigen::Index iter,
                                std::string_view failed_solver,
                                std::string_view next_solver,
                                const std::exception& e) {
  if (!msgs) {
    return;
  }
  // Build once so we don't interleave with other logs.
  std::ostringstream os;
  os << "[" << context << "] WARNING: solver fallback\n"
     << "  " << std::left << std::setw(12) << "iteration:" << iter << "\n"
     << "  " << std::left << std::setw(12) << "failed:" << failed_solver << "\n"
     << "  " << std::left << std::setw(12) << "reason:" << e.what() << "\n"
     << "  " << std::left << std::setw(12) << "action:"
     << "trying " << next_solver << "\n"
     << "note: this warning message will only be displayed once."
     << "\n";
  (*msgs) << os.str();
}
 
template <bool InitTheta, typename Opts>
inline decltype(auto) theta_init_impl(Eigen::Index theta_size, Opts&& options) {
  if constexpr (InitTheta) {
    // If requested, use the prior mean as the initial value
    return std::decay_t<decltype(options)>(options).theta_0;
  } else {
    return Eigen::MatrixXd::Zero(theta_size, 1);
  }
}
 
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_);
      };
}
 
static STAN_THREADS_DEF std::once_flag fallback_warning_1_2;
static STAN_THREADS_DEF std::once_flag fallback_warning_2_3;
template <typename LLFun, typename LLTupleArgs, typename CovarMat,
          bool InitTheta,
          require_t<is_all_arithmetic_scalar<CovarMat, LLTupleArgs>>* = nullptr>
inline auto laplace_marginal_density_est(
    LLFun&& ll_fun, LLTupleArgs&& ll_args, CovarMat&& covariance,
    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 {
    return -0.5 * a_val.dot(theta_val)
           + laplace_likelihood::log_likelihood(ll_fun, theta_val, ll_args,
                                                msgs);
  };
  auto theta_grad_f = [&ll_fun, &ll_args, &msgs](auto&& theta_val) {
    return laplace_likelihood::theta_grad(ll_fun, theta_val, ll_args, msgs);
  };
  decltype(auto) theta_init = theta_init_impl<InitTheta>(theta_size, options);
  // When the user supplies a non-zero theta_init, we must initialise a
  // consistently so that the invariant theta = Sigma * a holds.  Otherwise
  // the prior term -0.5 * a'*theta vanishes (a=0 while theta!=0), inflating
  // the initial objective and causing the Wolfe line search to reject the
  // first Newton step.
  auto state
      = NewtonState(theta_size, obj_fun, theta_grad_f, covariance, theta_init);
  // 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) {
      if (options.hessian_block_size == 1) {
        CholeskyWSolverDiag solver(state, covariance);
        return run_newton_loop(solver, state, options, step_iter, ll_fun,
                               ll_args, covariance, update_fun, msgs);
      } else {
        CholeskyWSolverBlock solver(state, options.hessian_block_size);
        return run_newton_loop(solver, state, options, step_iter, ll_fun,
                               ll_args, covariance, update_fun, msgs);
      }
    }
  } catch (const std::exception& e) {
    const std::string solver_type
        = (options.hessian_block_size == 1) ? "Diagonal" : "Block";
    std::string failed = "solver 1 (" + solver_type + " Hessian-root Cholesky)";
    if (!options.allow_fallthrough) {
      throw_solver_failure("laplace_marginal_density", step_iter, failed, e);
    }
    std::call_once(
        fallback_warning_1_2,
        [](auto&&... args) STAN_COLD_PATH {
          log_solver_fallback(std::forward<decltype(args)>(args)...);
        },
        msgs, "laplace_marginal_density", step_iter, std::move(failed),
        "solver 2 (Covariance-root Cholesky)", e);
  }
  try {
    if (options.solver == 2 || options.allow_fallthrough) {
      CholeskyKSolver solver(state, covariance);
      return run_newton_loop(solver, state, options, step_iter, ll_fun, ll_args,
                             covariance, update_fun, msgs);
    }
  } catch (const std::exception& e) {
    if (!options.allow_fallthrough) {
      throw_solver_failure("laplace_marginal_density", step_iter,
                           "solver 2 (Covariance-root Cholesky)", e);
    }
    std::call_once(
        fallback_warning_2_3,
        [](auto&&... args) STAN_COLD_PATH {
          log_solver_fallback(std::forward<decltype(args)>(args)...);
        },
        msgs, "laplace_marginal_density", step_iter,
        "solver 2 (Covariance-root Cholesky)", "solver 3 (General LU solver)",
        e);
  }
  if (options.solver == 3 || options.allow_fallthrough) {
    LUSolver solver;
    return run_newton_loop(solver, state, options, step_iter, ll_fun, ll_args,
                           covariance, update_fun, msgs);
  }
  throw std::domain_error(
      std::string("You chose a solver (") + std::to_string(options.solver)
      + ") that is not valid. Please choose either 1, 2, or 3.");
}
}  // namespace internal
}  // namespace math
}  // namespace stan
#endif