[WIP] Imbedded Laplace Approximation

stan/math/mix.hpp

@@ -26,4 +26,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/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,54 @@
+#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 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 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();
+    v(Eigen::seq(i, x_size - 1, hessian_block_size)).setOnes();
+    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,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,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,91 @@
+#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/fun.hpp>
+#include <stan/math/prim/prob/multi_normal_cholesky_rng.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 Laplace approximation of p(theta_pred | y, phi, x_pred).
+ * Note that while the data is observed at x (train_tuple), the new samples
+ * are drawn for covariates x_pred (pred_tuple).
+ * To sample the "original" theta's, set pred_tuple = train_tuple.
+ * @tparam D Type of likelihood function.
+ * @tparam LLArgs Type of arguments of likelihood function.
+ * @tparam ThetaMatrix Type of latent Gaussian variables.
+ * @tparam EtaMatrix Type of additional arguments for likelihood function.
+ * @tparam CovarFun Type of covariance function.
+ * @tparam RNG Type of RNG number.
+ * @tparam TrainTuple Type of observed/training covariate.
+ * @tparam PredTuple Type of predictive covariate.
+ * @tparam CovarArgs Tuple of additional arguments for the covariance function
+ * @param ll_fun Likelihood function.
+ * @param ll_args Arguments for likelihood function.
+ * @param covariance_function Covariance function.
+ * @param eta Additional arguments for likelihood function.
+ * @param theta_0 Initial guess for finding the mode of the conditional
+                  pi(theta_pred | y, phi, x_pred).
+ * @param options Control parameter for optimizer underlying Laplace approx.
+ * @param train_tuple Observed/training covariates for covariance function.
+ * @param pred_tuple Predictive covariates for covariance function.
+ * @param rng Rng number.
+ * @param msgs Stream for function prints.
+ * @param args Variadic arguments for likelihood function.
+ */
+template <typename D, typename LLArgs, typename ThetaMatrix, typename EtaMatrix,
+          typename CovarFun, class RNG, typename TrainTuple, typename PredTuple,
+          typename CovarArgs,
+          require_all_eigen_t<ThetaMatrix, EtaMatrix>* = nullptr>
+inline Eigen::VectorXd laplace_base_rng(
+    D&& ll_fun, LLArgs&& ll_args, CovarFun&& covariance_function,
+    const EtaMatrix& eta, const ThetaMatrix& theta_0,
+    const laplace_options& options, TrainTuple&& train_tuple,
+    PredTuple&& pred_tuple, RNG& rng, std::ostream* msgs,
+    CovarArgs&& covar_args) {
+  using Eigen::MatrixXd;
+  using Eigen::VectorXd;
+  auto covar_args_val = stan::math::to_ref(value_of(covar_args));
+  auto eta_dbl = value_of(eta);
+  auto md_est = laplace_marginal_density_est(
+      ll_fun, ll_args, eta_dbl, value_of(theta_0), covariance_function,
+      std::tuple_cat(std::forward<TrainTuple>(train_tuple), covar_args_val),
+      options, msgs);
+  // 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), covar_args_val));
+  VectorXd pred_mean = covariance_pred * md_est.theta_grad;
+  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