[WIP] Embedded Laplace Approximation

.gitignore

@@ -35,6 +35,8 @@ benchmarks/*.csv
 *.o-*
 *.exe
 *.a
+# mac debug folders
+*.dSYM/
 
 # Intel template building blocks (TBB)
 lib/tbb

Jenkinsfile

@@ -322,6 +322,31 @@ pipeline {
                     }
                     post { always { retry(3) { deleteDir() } } }
                 }
+                stage('Laplace Unit Tests') {
+                    agent {
+                        docker {
+                            image 'stanorg/ci:gpu-cpp17'
+                            label 'linux'
+                            args '--cap-add SYS_PTRACE'
+                        }
+                    }
+                    when {
+                        expression {
+                            !skipRemainingStages
+                        }
+                    }
+                    steps {
+                        unstash 'MathSetup'
+                        sh "echo CXXFLAGS += -fsanitize=address -march=native -mtune=native >> make/local"
+                        script {
+                            if (!(params.optimizeUnitTests || isBranch('develop') || isBranch('master'))) {
+                                sh "echo O=3 >> make/local"
+                            }
+                            runTests("test/unit/math/laplace/*_test.cpp", false)
+                        }
+                    }
+                    post { always { retry(3) { deleteDir() } } }
+                }
                 stage('OpenCL GPU tests') {
                     agent {
                         docker {

stan/math/fwd/meta/is_fvar.hpp

@@ -21,5 +21,8 @@ struct is_fvar<T,
                std::enable_if_t<internal::is_fvar_impl<std::decay_t<T>>::value>>
     : std::true_type {};
 
+template <typename T>
+inline constexpr bool is_fvar_v = is_fvar<T>::value;
+
 }  // namespace stan
 #endif

stan/math/mix.hpp

@@ -1,6 +1,10 @@
 #ifndef STAN_MATH_MIX_HPP
 #define STAN_MATH_MIX_HPP
 
+#include <stan/math/fwd/fun/Eigen_NumTraits.hpp>
+#include <stan/math/rev/core/Eigen_NumTraits.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+
 #include <stan/math/mix/meta.hpp>
 #include <stan/math/mix/fun.hpp>
 #include <stan/math/mix/functor.hpp>
@@ -26,4 +30,6 @@
 
 #include <stan/math/prim.hpp>
 
+#include <stan/math/mix/prob.hpp>
+
 #endif

stan/math/mix/functor.hpp

@@ -2,13 +2,16 @@
 #define STAN_MATH_MIX_FUNCTOR_HPP
 
 #include <stan/math/mix/functor/derivative.hpp>
-#include <stan/math/mix/functor/finite_diff_grad_hessian.hpp>
 #include <stan/math/mix/functor/finite_diff_grad_hessian_auto.hpp>
+#include <stan/math/mix/functor/finite_diff_grad_hessian.hpp>
 #include <stan/math/mix/functor/grad_hessian.hpp>
 #include <stan/math/mix/functor/grad_tr_mat_times_hessian.hpp>
 #include <stan/math/mix/functor/gradient_dot_vector.hpp>
 #include <stan/math/mix/functor/hessian.hpp>
+#include <stan/math/mix/functor/laplace_base_rng.hpp>
+#include <stan/math/mix/functor/laplace_likelihood.hpp>
+#include <stan/math/mix/functor/laplace_marginal_density.hpp>
+#include <stan/math/mix/functor/hessian_block_diag.hpp>
 #include <stan/math/mix/functor/hessian_times_vector.hpp>
 #include <stan/math/mix/functor/partial_derivative.hpp>
-
 #endif

stan/math/mix/functor/derivative.hpp

@@ -1,9 +1,9 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_DERIVATIVE_HPP
 #define STAN_MATH_MIX_FUNCTOR_DERIVATIVE_HPP
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/fwd/core.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <vector>
 
 namespace stan {
@@ -21,7 +21,7 @@ namespace math {
  * @param[out] dfx_dx Value of derivative
  */
 template <typename T, typename F>
-void derivative(const F& f, const T& x, T& fx, T& dfx_dx) {
+inline void derivative(const F& f, const T& x, T& fx, T& dfx_dx) {
   fvar<T> x_fvar = fvar<T>(x, 1.0);
   fvar<T> fx_fvar = f(x_fvar);
   fx = fx_fvar.val_;

stan/math/mix/functor/finite_diff_grad_hessian.hpp

@@ -38,10 +38,10 @@ namespace math {
  * @param[in] epsilon perturbation size
  */
 template <typename F>
-void finite_diff_grad_hessian(const F& f, const Eigen::VectorXd& x, double& fx,
-                              Eigen::MatrixXd& hess,
-                              std::vector<Eigen::MatrixXd>& grad_hess_fx,
-                              double epsilon = 1e-04) {
+inline void finite_diff_grad_hessian(const F& f, const Eigen::VectorXd& x,
+                                     double& fx, Eigen::MatrixXd& hess,
+                                     std::vector<Eigen::MatrixXd>& grad_hess_fx,
+                                     double epsilon = 1e-04) {
   int d = x.size();
   grad_hess_fx.clear();
 

stan/math/mix/functor/finite_diff_grad_hessian_auto.hpp

@@ -41,9 +41,9 @@ namespace math {
  * @param[out] grad_hess_fx gradient of Hessian of function at argument
  */
 template <typename F>
-void finite_diff_grad_hessian_auto(const F& f, const Eigen::VectorXd& x,
-                                   double& fx, Eigen::MatrixXd& hess,
-                                   std::vector<Eigen::MatrixXd>& grad_hess_fx) {
+inline void finite_diff_grad_hessian_auto(
+    const F& f, const Eigen::VectorXd& x, double& fx, Eigen::MatrixXd& hess,
+    std::vector<Eigen::MatrixXd>& grad_hess_fx) {
   int d = x.size();
 
   grad_hess_fx.clear();

stan/math/mix/functor/grad_hessian.hpp

@@ -1,9 +1,9 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_GRAD_HESSIAN_HPP
 #define STAN_MATH_MIX_FUNCTOR_GRAD_HESSIAN_HPP
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/fwd/core.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <stdexcept>
 #include <vector>
 
@@ -39,7 +39,7 @@ namespace math {
  * @param[out] grad_H Gradient of the Hessian of function at argument
  */
 template <typename F>
-void grad_hessian(
+inline void grad_hessian(
     const F& f, const Eigen::Matrix<double, Eigen::Dynamic, 1>& x, double& fx,
     Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& H,
     std::vector<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> >&

stan/math/mix/functor/grad_tr_mat_times_hessian.hpp

@@ -1,18 +1,18 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_GRAD_TR_MAT_TIMES_HESSIAN_HPP
 #define STAN_MATH_MIX_FUNCTOR_GRAD_TR_MAT_TIMES_HESSIAN_HPP
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/mix/functor/gradient_dot_vector.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/fwd/core.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <stdexcept>
 #include <vector>
 
 namespace stan {
 namespace math {
 
 template <typename F>
-void grad_tr_mat_times_hessian(
+inline void grad_tr_mat_times_hessian(
     const F& f, const Eigen::Matrix<double, Eigen::Dynamic, 1>& x,
     const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& M,
     Eigen::Matrix<double, Eigen::Dynamic, 1>& grad_tr_MH) {
@@ -26,7 +26,7 @@ void grad_tr_mat_times_hessian(
 
   Matrix<var, Dynamic, 1> x_var(x.size());
   for (int i = 0; i < x.size(); ++i) {
-    x_var(i) = x(i);
+    x_var.coeffRef(i) = x(i);
   }
 
   Matrix<fvar<var>, Dynamic, 1> x_fvar(x.size());

stan/math/mix/functor/gradient_dot_vector.hpp

@@ -1,19 +1,19 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_GRADIENT_DOT_VECTOR_HPP
 #define STAN_MATH_MIX_FUNCTOR_GRADIENT_DOT_VECTOR_HPP
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/fwd/core.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <vector>
 
 namespace stan {
 namespace math {
 
 template <typename T1, typename T2, typename F>
-void gradient_dot_vector(const F& f,
-                         const Eigen::Matrix<T1, Eigen::Dynamic, 1>& x,
-                         const Eigen::Matrix<T2, Eigen::Dynamic, 1>& v, T1& fx,
-                         T1& grad_fx_dot_v) {
+inline void gradient_dot_vector(const F& f,
+                                const Eigen::Matrix<T1, Eigen::Dynamic, 1>& x,
+                                const Eigen::Matrix<T2, Eigen::Dynamic, 1>& v,
+                                T1& fx, T1& grad_fx_dot_v) {
   using Eigen::Matrix;
   Matrix<fvar<T1>, Eigen::Dynamic, 1> x_fvar(x.size());
   for (int i = 0; i < x.size(); ++i) {

stan/math/mix/functor/hessian.hpp

@@ -1,9 +1,9 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_HESSIAN_HPP
 #define STAN_MATH_MIX_FUNCTOR_HESSIAN_HPP
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/fwd/core.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <stdexcept>
 
 namespace stan {
@@ -39,9 +39,10 @@ namespace math {
  * @param[out] H Hessian of function at argument
  */
 template <typename F>
-void hessian(const F& f, const Eigen::Matrix<double, Eigen::Dynamic, 1>& x,
-             double& fx, Eigen::Matrix<double, Eigen::Dynamic, 1>& grad,
-             Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& H) {
+inline void hessian(const F& f,
+                    const Eigen::Matrix<double, Eigen::Dynamic, 1>& x,
+                    double& fx, Eigen::Matrix<double, Eigen::Dynamic, 1>& grad,
+                    Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>& H) {
   H.resize(x.size(), x.size());
   grad.resize(x.size());
 

stan/math/mix/functor/hessian_block_diag.hpp

@@ -0,0 +1,56 @@
+#ifndef STAN_MATH_MIX_FUNCTOR_HESSIAN_BLOCK_DIAG_HPP
+#define STAN_MATH_MIX_FUNCTOR_HESSIAN_BLOCK_DIAG_HPP
+
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/mix/functor/hessian_times_vector.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Returns a block diagonal Hessian by computing the relevant directional
+ * derivatives and storing them in a matrix.
+ * For m the size of each block, the operations const m calls to
+ * hessian_times_vector, that is m forward sweeps and m reverse sweeps.
+ * @tparam F Type of function to differentiate.
+ * @tparam Eta Type of additional arguments passed to F.
+ * @tparam Args Type of variadic arguments passed to F.
+ * @param f Function to differentiate.
+ * @param x Arguments with respect to which we differentiate.
+ * @param eta Additional arguments for f.
+ * @param hessian_block_size
+ * @param args Additional variadic arguments for f.
+ */
+template <typename F, typename... Args>
+inline Eigen::SparseMatrix<double> hessian_block_diag(
+    F&& f, const Eigen::VectorXd& x, const Eigen::Index hessian_block_size,
+    Args&&... args) {
+  using Eigen::MatrixXd;
+  using Eigen::VectorXd;
+
+  const Eigen::Index x_size = x.size();
+  Eigen::SparseMatrix<double> H(x_size, x_size);
+  H.reserve(Eigen::VectorXi::Constant(x_size, hessian_block_size));
+  VectorXd v(x_size);
+  Eigen::Index n_blocks = x_size / hessian_block_size;
+  Eigen::VectorXd Hv = Eigen::VectorXd::Zero(x_size);
+  for (Eigen::Index i = 0; i < hessian_block_size; ++i) {
+    v.setZero();
+    for (Eigen::Index j = i; j < x_size; j += hessian_block_size) {
+      v.coeffRef(j) = 1;
+    }
+    hessian_times_vector(f, Hv, x, v, args...);
+    for (int j = 0; j < n_blocks; ++j) {
+      for (int k = 0; k < hessian_block_size; ++k) {
+        H.insert(k + j * hessian_block_size, i + j * hessian_block_size)
+            = Hv(k + j * hessian_block_size);
+      }
+    }
+  }
+  return H;
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/mix/functor/hessian_times_vector.hpp

@@ -1,21 +1,20 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_HESSIAN_TIMES_VECTOR_HPP
 #define STAN_MATH_MIX_FUNCTOR_HESSIAN_TIMES_VECTOR_HPP
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/fwd/core.hpp>
-#include <stan/math/prim/fun/Eigen.hpp>
 #include <stdexcept>
 #include <vector>
 
 namespace stan {
 namespace math {
 
 template <typename F>
-void hessian_times_vector(const F& f,
-                          const Eigen::Matrix<double, Eigen::Dynamic, 1>& x,
-                          const Eigen::Matrix<double, Eigen::Dynamic, 1>& v,
-                          double& fx,
-                          Eigen::Matrix<double, Eigen::Dynamic, 1>& Hv) {
+inline void hessian_times_vector(
+    const F& f, const Eigen::Matrix<double, Eigen::Dynamic, 1>& x,
+    const Eigen::Matrix<double, Eigen::Dynamic, 1>& v, double& fx,
+    Eigen::Matrix<double, Eigen::Dynamic, 1>& Hv) {
   using Eigen::Matrix;
 
   // Run nested autodiff in this scope
@@ -35,10 +34,13 @@ void hessian_times_vector(const F& f,
     Hv(i) = x_var(i).adj();
   }
 }
-template <typename T, typename F>
+
+template <typename T, typename F, typename EigVec,
+          require_eigen_vector_t<EigVec>* = nullptr,
+          require_stan_scalar_t<T>* = nullptr>
 void hessian_times_vector(const F& f,
                           const Eigen::Matrix<T, Eigen::Dynamic, 1>& x,
-                          const Eigen::Matrix<T, Eigen::Dynamic, 1>& v, T& fx,
+                          const EigVec& v, T& fx,
                           Eigen::Matrix<T, Eigen::Dynamic, 1>& Hv) {
   using Eigen::Matrix;
   Matrix<T, Eigen::Dynamic, 1> grad;
@@ -47,6 +49,29 @@ void hessian_times_vector(const F& f,
   Hv = H * v;
 }
 
+/**
+ * Overload Hessian_times_vector function, under stan/math/mix/functor
+ * to handle functions which take in arguments
+ * and pstream.
+ */
+template <typename F, typename XAdj, typename XVec, typename VVec,
+          typename... Args,
+          require_all_eigen_vector_t<XAdj, XVec, VVec>* = nullptr>
+inline void hessian_times_vector(const F& f, XAdj& x_adj, XVec&& x, VVec&& v,
+                                 Args&&... args) {
+  nested_rev_autodiff nested;
+  const Eigen::Index x_size = x.size();
+  Eigen::Matrix<var, Eigen::Dynamic, 1> x_var = std::forward<XVec>(x);
+  Eigen::Matrix<fvar<var>, Eigen::Dynamic, 1> x_fvar(x_size);
+  for (Eigen::Index i = 0; i < x_size; i++) {
+    x_fvar(i) = fvar<var>(x_var(i), v(i));
+  }
+  fvar<var> fx_fvar = f(x_fvar, args...);
+  grad(fx_fvar.d_.vi_);
+  x_adj = x_var.adj();
+}
+
 }  // namespace math
 }  // namespace stan
+
 #endif