[WIP] Imbedded Laplace Approximation

make/compiler_flags

@@ -120,32 +120,9 @@ INC_GTEST ?= -I $(GTEST)/include -I $(GTEST)
 CPPFLAGS_BOOST ?= -DBOOST_DISABLE_ASSERTS
 CPPFLAGS_SUNDIALS ?= -DNO_FPRINTF_OUTPUT $(CPPFLAGS_OPTIM_SUNDIALS) $(CXXFLAGS_FLTO_SUNDIALS)
 #CPPFLAGS_GTEST ?=
-STAN_HAS_CXX17 ?= false
-ifeq ($(CXX_TYPE), gcc)
-  GCC_GE_73 := $(shell [ $(CXX_MAJOR) -gt 7 -o \( $(CXX_MAJOR) -eq 7 -a $(CXX_MINOR) -ge 1 \) ] && echo true)
-  ifeq ($(GCC_GE_73),true)
-    STAN_HAS_CXX17 := true
-  endif
-else ifeq ($(CXX_TYPE), clang)
-  CLANG_GE_5 := $(shell [ $(CXX_MAJOR) -gt 5 -o \( $(CXX_MAJOR) -eq 5 -a $(CXX_MINOR) -ge 0 \) ] && echo true)
-  ifeq ($(CLANG_GE_5),true)
-    STAN_HAS_CXX17 := true
-  endif
-else ifeq ($(CXX_TYPE), mingw32-gcc)
-  MINGW_GE_50 := $(shell [ $(CXX_MAJOR) -gt 5 -o \( $(CXX_MAJOR) -eq 5 -a $(CXX_MINOR) -ge 0 \) ] && echo true)
-  ifeq ($(MINGW_GE_50),true)
-    STAN_HAS_CXX17 := true
-  endif
-endif
 
-ifeq ($(STAN_HAS_CXX17), true)
-  CXXFLAGS_LANG ?= -std=c++17
-  CXXFLAGS_STANDARD ?= c++17
-else
-  $(warning "Stan cannot detect if your compiler has the C++17 standard. If it does, please set STAN_HAS_CXX17=true in your make/local file. C++17 support is mandatory in the next release of Stan. Defaulting to C++14")
-  CXXFLAGS_LANG ?= -std=c++1y
-  CXXFLAGS_STANDARD ?= c++1y
-endif
+CXXFLAGS_LANG ?= -std=c++17
+CXXFLAGS_STANDARD ?= c++17
 #CXXFLAGS_BOOST ?=
 CXXFLAGS_SUNDIALS ?= -pipe $(CXXFLAGS_OPTIM_SUNDIALS) $(CPPFLAGS_FLTO_SUNDIALS)
 #CXXFLAGS_GTEST

stan/math/mix.hpp

@@ -1,10 +1,5 @@
 #ifndef STAN_MATH_MIX_HPP
 #define STAN_MATH_MIX_HPP
-
-#include <stan/math/mix/meta.hpp>
-#include <stan/math/mix/fun.hpp>
-#include <stan/math/mix/functor.hpp>
-
 #include <stan/math/fwd/constraint.hpp>
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/meta.hpp>
@@ -26,4 +21,9 @@
 
 #include <stan/math/prim.hpp>
 
+#include <stan/math/mix/meta.hpp>
+#include <stan/math/mix/fun.hpp>
+#include <stan/math/mix/functor.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/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/mix/functor/hessian_times_vector.hpp>
+#include <Eigen/Sparse>
+
+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
+ * @tparam Eta
+ * @tparam Args
+ * @param f
+ * @param x
+ * @param eta
+ * @param hessian_block_size
+ * @param args
+ */
+template <typename F, typename Eta, typename... Args,
+          require_eigen_t<Eta>* = nullptr>
+inline Eigen::SparseMatrix<double> hessian_block_diag(
+    F&& f, const Eigen::VectorXd& x, const Eta& eta,
+    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;
+  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(j) = 1;
+    }
+    VectorXd Hv = hessian_times_vector(f, x, eta, 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,11 +1,10 @@
 #ifndef STAN_MATH_MIX_FUNCTOR_HESSIAN_TIMES_VECTOR_HPP
 #define STAN_MATH_MIX_FUNCTOR_HESSIAN_TIMES_VECTOR_HPP
 
-#include <stan/math/fwd/core.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/fwd/core.hpp>
 #include <stan/math/rev/core.hpp>
-#include <stdexcept>
-#include <vector>
+#include <stan/math/prim/meta.hpp>
 
 namespace stan {
 namespace math {
@@ -35,6 +34,7 @@ void hessian_times_vector(const F& f,
     Hv(i) = x_var(i).adj();
   }
 }
+
 template <typename T, typename F>
 void hessian_times_vector(const F& f,
                           const Eigen::Matrix<T, Eigen::Dynamic, 1>& x,
@@ -47,6 +47,31 @@ 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 eta, delta, delta_int,
+ * and pstream.
+ */
+template <typename F, typename Eta, require_eigen_t<Eta>* = nullptr,
+          typename... Args>
+inline Eigen::VectorXd hessian_times_vector(const F& f,
+                                            const Eigen::VectorXd& x,
+                                            const Eta& eta,
+                                            const Eigen::VectorXd& v,
+                                            Args&&... args) {
+  nested_rev_autodiff nested;
+  const Eigen::Index x_size = x.size();
+  Eigen::Matrix<var, Eigen::Dynamic, 1> x_var = 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, eta, args...);
+  grad(fx_fvar.d_.vi_);
+  return x_var.adj();
+}
+
 }  // namespace math
 }  // namespace stan
+
 #endif

stan/math/mix/functor/laplace_base_rng.hpp

@@ -0,0 +1,92 @@
+#ifndef STAN_MATH_MIX_FUNCTOR_LAPLACE_BASE_RNG_HPP
+#define STAN_MATH_MIX_FUNCTOR_LAPLACE_BASE_RNG_HPP
+
+#include <stan/math/mix/functor/laplace_marginal_density.hpp>
+#include <stan/math/prim/prob/multi_normal_cholesky_rng.hpp>
+#include <stan/math/prim/fun.hpp>
+
+#include <Eigen/Sparse>
+
+namespace stan {
+namespace math {
+
+/**
+ * In a latent gaussian model,
+ *
+ *   theta ~ Normal(theta | 0, Sigma(phi, x))
+ *   y ~ pi(y | theta, eta)
+ *
+ * returns a multivariate normal random variate sampled
+ * from the gaussian approximation of p(theta_pred | y, phi, x_pred).
+ * Note that while the data is observed at x, the new samples
+ * are drawn for covariates x_pred.
+ * To sample the "original" theta's, set x_pred = x.
+ * @tparam D
+ * @tparam LLArgs
+ * @tparam ThetaMatrix
+ * @tparam EtaMatrix
+ * @tparam CovarFun
+ * @tparam RNG
+ * @tparam TrainTuple
+ * @tparam PredTuple
+ * @tparam Args
+ * @param ll_fun
+ * @param ll_args
+ * @param covariance_function
+ * @param eta
+ * @param theta_0
+ * @param options
+ * @param train_tuple
+ * @param pred_tuple
+ * @param rng
+ * @param msgs
+ * @param args
+ */
+template <typename D, typename LLArgs, typename ThetaMatrix, typename EtaMatrix,
+          typename CovarFun, class RNG, typename TrainTuple, typename PredTuple,
+          typename... Args,
+          require_all_eigen_t<ThetaMatrix, EtaMatrix>* = nullptr>
+inline Eigen::VectorXd laplace_base_rng(
+    D&& ll_fun, LLArgs&& ll_args, CovarFun&& covariance_function,
+    const ThetaMatrix& eta, const EtaMatrix& theta_0,
+    const laplace_options& options, TrainTuple&& train_tuple,
+    PredTuple&& pred_tuple, RNG& rng, std::ostream* msgs, Args&&... args) {
+  using Eigen::MatrixXd;
+  using Eigen::VectorXd;
+  auto args_dbl = std::make_tuple(to_ref(value_of(args))...);
+  auto eta_dbl = value_of(eta);
+  auto md_est = apply(
+      [&](auto&&... args_val) {
+        return laplace_marginal_density_est(
+            ll_fun, ll_args, covariance_function, eta_dbl, value_of(theta_0),
+            msgs, options, args_val...);
+      },
+      std::tuple_cat(std::forward<TrainTuple>(train_tuple), args_dbl));
+  // Modified R&W method
+  MatrixXd covariance_pred = apply(
+      [&covariance_function, &msgs](auto&&... args_val) {
+        return covariance_function(args_val..., msgs);
+      },
+      std::tuple_cat(std::forward<PredTuple>(pred_tuple), args_dbl));
+  VectorXd pred_mean = covariance_pred * md_est.l_grad.head(theta_0.rows());
+  if (options.solver == 1 || options.solver == 2) {
+    Eigen::MatrixXd V_dec = mdivide_left_tri<Eigen::Lower>(
+        md_est.L, md_est.W_r * covariance_pred);
+    Eigen::MatrixXd Sigma = covariance_pred - V_dec.transpose() * V_dec;
+    return multi_normal_rng(pred_mean, Sigma, rng);
+  } else {
+    Eigen::MatrixXd Sigma
+        = covariance_pred
+          - covariance_pred
+                * (md_est.W_r
+                   - md_est.W_r
+                         * md_est.LU.solve(md_est.covariance * md_est.W_r))
+                * covariance_pred;
+    return multi_normal_rng(pred_mean, Sigma, rng);
+  }
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif