Add Wolfe line search to Laplace approximation

stan/math/mix/functor/finite_diff_grad_hessian_auto.hpp

@@ -58,7 +58,7 @@ inline void finite_diff_grad_hessian_auto(
 
   for (int i = 0; i < d; ++i) {
     double dummy_fx_eval;
-    double epsilon = finite_diff_stepsize(x(i));
+    double epsilon = finite_diff_stepsize<2>(x(i));
     hess_diff.setZero();
 
     x_temp(i) = x(i) + 2 * epsilon;

stan/math/mix/functor/laplace_base_rng.hpp

@@ -38,12 +38,13 @@ inline Eigen::VectorXd laplace_base_rng(
     LLFunc&& ll_fun, LLArgs&& ll_args, CovarFun&& covariance_function,
     CovarArgs&& covar_args, const laplace_options<InitTheta>& options, RNG& rng,
     std::ostream* msgs) {
+  Eigen::MatrixXd covariance_train = stan::math::apply(
+      [msgs, &covariance_function](auto&&... args) {
+        return covariance_function(std::forward<decltype(args)>(args)..., msgs);
+      },
+      std::forward<CovarArgs>(covar_args));
   auto md_est = internal::laplace_marginal_density_est(
-      ll_fun, std::forward<LLArgs>(ll_args),
-      std::forward<CovarFun>(covariance_function),
-      to_ref(std::forward<CovarArgs>(covar_args)), options, msgs);
-  // Modified R&W method
-  auto&& covariance_train = md_est.covariance;
+      ll_fun, std::forward<LLArgs>(ll_args), covariance_train, options, msgs);
   Eigen::VectorXd mean_train = covariance_train * md_est.theta_grad;
   if (options.solver == 1 || options.solver == 2) {
     Eigen::MatrixXd V_dec

stan/math/mix/functor/laplace_likelihood.hpp

@@ -8,11 +8,49 @@
 namespace stan {
 namespace math {
 
+namespace internal {
+/**
+ * Set all adjoints of the output to zero.
+ */
+template <typename Output>
+inline void set_zero_adjoint(Output&& output) {
+  if constexpr (is_all_arithmetic_scalar_v<Output>) {
+    return;
+  } else {
+    return iter_tuple_nested(
+        [](auto&& output_i) {
+          using output_i_t = std::decay_t<decltype(output_i)>;
+          if constexpr (is_all_arithmetic_scalar_v<output_i_t>) {
+            return;
+          } else if constexpr (is_std_vector<output_i_t>::value) {
+            for (Eigen::Index i = 0; i < output_i.size(); ++i) {
+              output_i[i].adj() = 0;
+            }
+          } else if constexpr (is_eigen_v<output_i_t>) {
+            output_i.adj().setZero();
+          } else if constexpr (is_stan_scalar_v<output_i_t>) {
+            output_i.adj() = 0;
+          } else {
+            static_assert(
+                sizeof(std::decay_t<output_i_t>*) == 0,
+                "INTERNAL ERROR:(laplace_marginal_lpdf) set_zero_adjoints was "
+                "not able to deduce the actions needed for the given type. "
+                "This is an internal error, please report it: "
+                "https://github.com/stan-dev/math/issues");
+          }
+        },
+        std::forward<Output>(output));
+  }
+}
+
+}  // namespace internal
+
 /**
  * functions to compute the log density, first, second,
  * and third-order derivatives for a likelihoood specified by the user.
  */
 namespace laplace_likelihood {
+
 namespace internal {
 /**
  * @tparam F A functor with `opertor()(Args&&...)` returning a scalar
@@ -106,6 +144,126 @@ inline auto shallow_copy_vargs(Args&&... args) {
       std::forward<Args>(args)...);
 }
 
+/**
+ * Computes theta gradient `f` wrt `theta` and `args...`
+ * @note If `Args` contains \ref var types then their adjoints will be
+ * calculated as a side effect.
+ * @tparam F A functor with `opertor()(Args&&...)` returning a scalar
+ * @tparam Theta A class assignable to an Eigen vector type
+ * @tparam Stream Type of stream for messages.
+ * @tparam Args Type of variadic arguments.
+ * @param f Log likelihood function.
+ * @param theta Latent Gaussian model.
+ * @param msgs Stream for messages.
+ * @param args Variadic arguments for the likelihood function.
+ */
+template <typename F, typename Theta, typename Stream, typename... Args,
+          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
+inline auto theta_grad(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
+  using Eigen::Dynamic;
+  using Eigen::Matrix;
+  nested_rev_autodiff nested;
+  Matrix<var, Dynamic, 1> theta_var = theta;
+  var f_var = f(theta_var, args..., msgs);
+  grad(f_var.vi_);
+  return theta_var.adj().eval();
+}
+
+/**
+ * Computes likelihood argument gradient of `f`
+ * @note If `Args` contains \ref var types then their adjoints will be
+ * calculated as a side effect.
+ * @tparam F A functor with `opertor()(Args&&...)` returning a scalar
+ * @tparam Theta A class assignable to an Eigen vector type
+ * @tparam Stream Type of stream for messages.
+ * @tparam Args Type of variadic arguments.
+ * @param f Log likelihood function.
+ * @param theta Latent Gaussian model.
+ * @param msgs Stream for messages.
+ * @param args Variadic arguments for the likelihood function.
+ */
+template <typename F, typename Theta, typename Stream, typename... Args,
+          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
+inline void ll_arg_grad(F&& f, Theta&& theta, Stream* msgs, Args&&... args) {
+  using Eigen::Dynamic;
+  using Eigen::Matrix;
+  nested_rev_autodiff nested;
+  var f_var = f(theta, args..., msgs);
+  grad(f_var.vi_);
+}
+
+/**
+ * Computes negative block diagonal Hessian of `f` wrt`theta` and `args...`
+ * @note If `Args` contains \ref var types then their adjoints will be
+ * calculated as a side effect.
+ * @tparam F A functor with `opertor()(Args&&...)` returning a scalar
+ * @tparam Theta A class assignable to an Eigen vector type
+ * @tparam Stream Type of stream for messages.
+ * @tparam Args Type of variadic arguments.
+ * @param f Log likelihood function.
+ * @param theta Latent Gaussian model.
+ * @param hessian_block_size If the Hessian of the log likelihood function w.r.t
+ *                           the latent Gaussian variable is block-diagonal,
+ *                           size of each block.
+ * @param msgs Stream for messages.
+ * @param args Variadic arguments for the likelihood function.
+ */
+template <typename F, typename Theta, typename Stream, typename... Args,
+          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
+inline auto diagonal_hessian(F&& f, Theta&& theta, Stream* msgs,
+                             Args&&... args) {
+  using Eigen::Dynamic;
+  using Eigen::Matrix;
+  const Eigen::Index theta_size = theta.size();
+  auto v = Eigen::VectorXd::Ones(theta_size);
+  Eigen::VectorXd hessian_v = Eigen::VectorXd::Zero(theta_size);
+  hessian_times_vector(f, hessian_v, std::forward<Theta>(theta), std::move(v),
+                       value_of(args)..., msgs);
+  return (-hessian_v).eval();
+}
+
+/**
+ * Computes negative block diagonal Hessian of `f` wrt`theta` and `args...`
+ * @note If `Args` contains \ref var types then their adjoints will be
+ * calculated as a side effect.
+ * @tparam F A functor with `opertor()(Args&&...)` returning a scalar
+ * @tparam Theta A class assignable to an Eigen vector type
+ * @tparam Stream Type of stream for messages.
+ * @tparam Args Type of variadic arguments.
+ * @param f Log likelihood function.
+ * @param theta Latent Gaussian model.
+ * @param hessian_block_size If the Hessian of the log likelihood function w.r.t
+ *                           the latent Gaussian variable is block-diagonal,
+ *                           size of each block.
+ * @param msgs Stream for messages.
+ * @param args Variadic arguments for the likelihood function.
+ */
+template <typename F, typename Theta, typename Stream, typename... Args,
+          require_eigen_vector_vt<std::is_arithmetic, Theta>* = nullptr>
+inline auto block_hessian(F&& f, Theta&& theta,
+                          const Eigen::Index hessian_block_size, Stream* msgs,
+                          Args&&... args) {
+  using Eigen::Dynamic;
+  using Eigen::Matrix;
+  const Eigen::Index theta_size = theta.size();
+  if (hessian_block_size == 1) {
+    auto v = Eigen::VectorXd::Ones(theta_size);
+    Eigen::VectorXd hessian_v = Eigen::VectorXd::Zero(theta_size);
+    hessian_times_vector(f, hessian_v, std::forward<Theta>(theta), std::move(v),
+                         value_of(args)..., msgs);
+    Eigen::SparseMatrix<double> hessian_theta(theta_size, theta_size);
+    hessian_theta.reserve(Eigen::VectorXi::Constant(theta_size, 1));
+    for (Eigen::Index i = 0; i < theta_size; i++) {
+      hessian_theta.insert(i, i) = hessian_v(i);
+    }
+    return (-hessian_theta).eval();
+  } else {
+    return (-hessian_block_diag(f, std::forward<Theta>(theta),
+                                hessian_block_size, value_of(args)..., msgs))
+        .eval();
+  }
+}
+
 /**
  * Computes theta gradient and negative block diagonal Hessian of `f` wrt
  * `theta` and `args...`
@@ -301,6 +459,79 @@ inline auto diff_eta_implicit(F&& f, V_t&& v, Theta&& theta, Stream* msgs,
 
 }  // namespace internal
 
+/**
+ * A wrapper that accepts a tuple as arguments.
+ * @tparam F A functor with `opertor()(Args&&...)` returning a scalar
+ * @tparam Theta A class assignable to an Eigen vector type
+ * @tparam TupleArgs Type of arguments for covariance function.
+ * @tparam Stream Type of stream for messages.
+ * @param f Log likelihood function.
+ * @param theta Latent Gaussian model.
+ * @param ll_tup Arguments for likelihood function
+ * @param msgs stream messages.
+ */
+template <typename F, typename Theta, typename TupleArgs, typename Stream,
+          require_eigen_vector_t<Theta>* = nullptr,
+          require_tuple_t<TupleArgs>* = nullptr>
+inline auto theta_grad(F&& f, Theta&& theta, TupleArgs&& ll_tup,
+                       Stream* msgs = nullptr) {
+  return apply(
+      [](auto&& f, auto&& theta, auto&& msgs, auto&&... args) {
+        return internal::theta_grad(std::forward<decltype(f)>(f),
+                                    std::forward<decltype(theta)>(theta), msgs,
+                                    std::forward<decltype(args)>(args)...);
+      },
+      std::forward<TupleArgs>(ll_tup), std::forward<F>(f),
+      std::forward<Theta>(theta), msgs);
+}
+
+template <typename F, typename Theta, typename TupleArgs, typename Stream,
+          require_eigen_vector_t<Theta>* = nullptr,
+          require_tuple_t<TupleArgs>* = nullptr>
+inline auto ll_arg_grad(F&& f, Theta&& theta, TupleArgs&& ll_tup,
+                        Stream* msgs = nullptr) {
+  return apply(
+      [](auto&& f, auto&& theta, auto&& msgs, auto&&... args) {
+        return internal::ll_arg_grad(std::forward<decltype(f)>(f),
+                                     std::forward<decltype(theta)>(theta), msgs,
+                                     std::forward<decltype(args)>(args)...);
+      },
+      std::forward<TupleArgs>(ll_tup), std::forward<F>(f),
+      std::forward<Theta>(theta), msgs);
+}
+
+template <typename F, typename Theta, typename TupleArgs, typename Stream,
+          require_eigen_vector_t<Theta>* = nullptr,
+          require_tuple_t<TupleArgs>* = nullptr>
+inline auto diagonal_hessian(F&& f, Theta&& theta, TupleArgs&& ll_tuple,
+                             Stream* msgs) {
+  return apply(
+      [](auto&& f, auto&& theta, auto* msgs, auto&&... args) {
+        return internal::diagonal_hessian(
+            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
+            msgs, std::forward<decltype(args)>(args)...);
+      },
+      std::forward<TupleArgs>(ll_tuple), std::forward<F>(f),
+      std::forward<Theta>(theta), msgs);
+}
+
+template <typename F, typename Theta, typename TupleArgs, typename Stream,
+          require_eigen_vector_t<Theta>* = nullptr,
+          require_tuple_t<TupleArgs>* = nullptr>
+inline auto block_hessian(F&& f, Theta&& theta,
+                          const Eigen::Index hessian_block_size,
+                          TupleArgs&& ll_tuple, Stream* msgs) {
+  return apply(
+      [](auto&& f, auto&& theta, auto hessian_block_size, auto* msgs,
+         auto&&... args) {
+        return internal::block_hessian(
+            std::forward<decltype(f)>(f), std::forward<decltype(theta)>(theta),
+            hessian_block_size, msgs, std::forward<decltype(args)>(args)...);
+      },
+      std::forward<TupleArgs>(ll_tuple), std::forward<F>(f),
+      std::forward<Theta>(theta), hessian_block_size, msgs);
+}
+
 /**
  * A wrapper that accepts a tuple as arguments.
  * @tparam F A functor with `opertor()(Args&&...)` returning a scalar