Merge pull request #3337 from florence-bockting/laplace/return-cov-chol

Add `laplace_latent_solve()` and `laplace_latent_tol_solve()`

Author: SteveBronder

stan/math/mix/functor/laplace_base_rng.hpp

@@ -5,6 +5,7 @@
 #include <stan/math/mix/functor/laplace_marginal_density.hpp>
 #include <stan/math/prim/prob/multi_normal_cholesky_rng.hpp>
 #include <stan/math/prim/prob/multi_normal_rng.hpp>
+#include <stan/math/prim/fun/cholesky_decompose.hpp>
 
 namespace stan {
 namespace math {
@@ -15,11 +16,19 @@ namespace math {
  *   theta ~ Normal(theta | 0, Sigma(phi, x))
  *   y ~ pi(y | theta, eta)
  *
- * returns a multivariate normal random variate sampled
+ * By default, returns a multivariate normal random variate sampled
  * from the Laplace approximation of p(theta_pred | y, phi, x_pred).
+ * If `ReturnMeanAndCovCholesky` is true, instead of drawing a sample this
+ * returns the posterior mean and the Cholesky factor of the posterior
+ * covariance of that same Laplace approximation, as a
+ * `std::tuple<Eigen::VectorXd, Eigen::MatrixXd>`.
  * 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 ReturnMeanAndCovCholesky If false (default), draw and return a
+ *   random variate from the approximate posterior. If true, return a tuple
+ *   containing the posterior mean and the lower-triangular Cholesky factor of
+ *   the posterior covariance instead of a sample.
  * @tparam LLFunc Type of likelihood function.
  * @tparam LLArgs Tuple of arguments types of likelihood function.
  * \laplace_common_template_args
@@ -31,13 +40,15 @@ namespace math {
  * \rng_arg
  * \msg_arg
  */
-template <typename LLFunc, typename LLArgs, typename CovarFun,
-          typename CovarArgs, bool InitTheta, typename RNG,
+template <bool ReturnMeanAndCovCholesky = false, typename LLFunc,
+          typename LLArgs, typename CovarFun, typename CovarArgs,
+          bool InitTheta, typename RNG,
           require_t<is_all_arithmetic_scalar<CovarArgs, LLArgs>>* = nullptr>
-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) {
+inline auto 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);
@@ -51,7 +62,12 @@ inline Eigen::VectorXd laplace_base_rng(
         = md_est.L.template triangularView<Eigen::Lower>().solve(
             md_est.W_r * covariance_train);
     Eigen::MatrixXd Sigma = covariance_train - V_dec.transpose() * V_dec;
-    return multi_normal_rng(std::move(mean_train), std::move(Sigma), rng);
+    if constexpr (ReturnMeanAndCovCholesky) {
+      Eigen::MatrixXd Sigma_chol = cholesky_decompose(Sigma);
+      return std::make_tuple(std::move(mean_train), std::move(Sigma_chol));
+    } else {
+      return multi_normal_rng(std::move(mean_train), std::move(Sigma), rng);
+    }
   } else {
     Eigen::MatrixXd Sigma
         = covariance_train
@@ -60,7 +76,12 @@ inline Eigen::VectorXd laplace_base_rng(
                    - md_est.W_r
                          * md_est.LU.solve(covariance_train * md_est.W_r))
                 * covariance_train;
-    return multi_normal_rng(std::move(mean_train), std::move(Sigma), rng);
+    if constexpr (ReturnMeanAndCovCholesky) {
+      Eigen::MatrixXd Sigma_chol = cholesky_decompose(Sigma);
+      return std::make_tuple(std::move(mean_train), std::move(Sigma_chol));
+    } else {
+      return multi_normal_rng(std::move(mean_train), std::move(Sigma), rng);
+    }
   }
 }
 

stan/math/mix/prob.hpp

@@ -5,6 +5,7 @@
 #include <stan/math/mix/prob/laplace_latent_poisson_log_rng.hpp>
 #include <stan/math/mix/prob/laplace_latent_neg_binomial_2_log_rng.hpp>
 #include <stan/math/mix/prob/laplace_latent_rng.hpp>
+#include <stan/math/mix/prob/laplace_latent_solve.hpp>
 #include <stan/math/mix/prob/laplace_marginal.hpp>
 #include <stan/math/mix/prob/laplace_marginal_neg_binomial_2_log_lpmf.hpp>
 #include <stan/math/mix/prob/laplace_marginal_bernoulli_logit_lpmf.hpp>

stan/math/mix/prob/laplace_latent_solve.hpp

@@ -0,0 +1,91 @@
+#ifndef STAN_MATH_MIX_PROB_LAPLACE_LATENT_SOLVE_HPP
+#define STAN_MATH_MIX_PROB_LAPLACE_LATENT_SOLVE_HPP
+
+#include <stan/math/mix/functor/laplace_base_rng.hpp>
+
+namespace stan {
+namespace math {
+
+namespace internal {
+// Placeholder RNG type: used whenever laplace_base_rng's
+// sampling branch is compiled out via `if constexpr`. In this case
+// no rng is needed as no sampling is performed.
+struct laplace_unused_rng {};
+
+}  // namespace internal
+
+/**
+ * In a latent gaussian model,
+ *
+ *   theta ~ Normal(0, Sigma(phi))
+ *   y ~ p(y|theta,phi)
+ *
+ * returns the posterior mean and Cholesky factor from the Laplace
+ * approximation to p(theta|y,phi), where the log likelihood is given by L_f.
+ * @tparam LLFunc Type of likelihood function.
+ * @tparam LLArgs Tuple of arguments types of likelihood function.
+ * \laplace_common_template_args
+ * @param ll_fun Likelihood function.
+ * @param ll_args Arguments for likelihood function.
+ * \laplace_common_args
+ * @param[in] hessian_block_size Block size for the Hessian approximation with
+ * respect to the latent gaussian variable theta.
+ * \laplace_options
+ * \msg_arg
+ */
+template <typename LLFunc, typename LLArgs, typename CovarFun,
+          typename CovarArgs, typename OpsTuple>
+inline auto laplace_latent_tol_solve(LLFunc&& ll_fun, LLArgs&& ll_args,
+                                     int hessian_block_size,
+                                     CovarFun&& covariance_function,
+                                     CovarArgs&& covar_args, OpsTuple&& ops,
+                                     std::ostream* msgs) {
+  auto options
+      = internal::tuple_to_laplace_options(std::forward<OpsTuple>(ops));
+  options.hessian_block_size = hessian_block_size;
+  internal::laplace_unused_rng unused_rng;
+  return laplace_base_rng<true>(std::forward<LLFunc>(ll_fun),
+                                std::forward<LLArgs>(ll_args),
+                                std::forward<CovarFun>(covariance_function),
+                                std::forward<CovarArgs>(covar_args),
+                                std::move(options), unused_rng, msgs);
+}
+
+/**
+ * In a latent gaussian model,
+ *
+ *   theta ~ Normal(0, Sigma(phi))
+ *   y ~ p(y|theta,phi)
+ *
+ * returns the posterior mean and Cholesky factor
+ * from the Laplace approximation of p(theta | y, phi).
+ * @tparam LLFunc Type of likelihood function.
+ * @tparam LLArgs Tuple of arguments types of likelihood function.
+ * \laplace_common_template_args
+ * @tparam RNG A valid boost rng type
+ * @param ll_fun Likelihood function.
+ * @param ll_args Arguments for likelihood function.
+ * \laplace_common_args
+ * @param[in] hessian_block_size Block size for the Hessian approximation with
+ * respect to the latent gaussian variable theta.
+ * \msg_arg
+ */
+template <typename LLFunc, typename LLArgs, typename CovarFun,
+          typename CovarArgs>
+inline auto laplace_latent_solve(LLFunc&& ll_fun, LLArgs&& ll_args,
+                                 int hessian_block_size,
+                                 CovarFun&& covariance_function,
+                                 CovarArgs&& covar_args, std::ostream* msgs) {
+  auto options = laplace_options_default{hessian_block_size};
+  internal::laplace_unused_rng unused_rng;
+  return laplace_base_rng<true>(std::forward<LLFunc>(ll_fun),
+                                std::forward<LLArgs>(ll_args),
+                                std::forward<CovarFun>(covariance_function),
+                                std::forward<CovarArgs>(covar_args),
+                                std::move(options), unused_rng, msgs);
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

test/unit/math/laplace/laplace_latent_solve_test.cpp

@@ -0,0 +1,76 @@
+#include <stan/math.hpp>
+#include <stan/math/mix.hpp>
+#include <test/unit/math/laplace/laplace_utility.hpp>
+
+#include <boost/random/mersenne_twister.hpp>
+
+#include <gtest/gtest.h>
+#include <stdexcept>
+#include <vector>
+
+namespace {
+struct poisson_log_likelihood {
+  template <typename Theta>
+  auto operator()(const Theta& theta, const std::vector<int>& y,
+                  std::ostream* pstream) const {
+    return stan::math::poisson_log_lpmf(y, theta);
+  }
+};
+}  // namespace
+
+TEST_F(laplace_count_two_dim_diag_test, latent_solve_mean_and_cov) {
+  using stan::math::laplace_latent_solve;
+  auto [mean_est, chol_est]
+      = laplace_latent_solve(poisson_log_likelihood{}, std::forward_as_tuple(y),
+                             1, stan::math::test::diagonal_kernel_functor{},
+                             std::forward_as_tuple(phi(0), phi(1)), nullptr);
+  constexpr double tol = 1e-6;
+  EXPECT_EQ(2, mean_est.size());
+  EXPECT_NEAR(theta_root(0), mean_est(0), tol);
+  EXPECT_NEAR(theta_root(1), mean_est(1), tol);
+  EXPECT_NEAR(0.0, chol_est(0, 1), 1e-12);  // check lower triangular matrix
+  Eigen::MatrixXd Sigma_est = chol_est * chol_est.transpose();
+  EXPECT_NEAR(K_laplace(0, 0), Sigma_est(0, 0), tol);
+  EXPECT_NEAR(K_laplace(1, 1), Sigma_est(1, 1), tol);
+  EXPECT_NEAR(K_laplace(0, 1), Sigma_est(0, 1), tol);
+  EXPECT_NEAR(K_laplace(1, 0), Sigma_est(1, 0), tol);
+}
+
+TEST_F(laplace_count_two_dim_diag_test, latent_tol_solve_mean_and_cov) {
+  using stan::math::laplace_latent_tol_solve;
+  constexpr double tolerance = 1e-12;
+  constexpr int max_num_steps = 1000;
+  constexpr int hessian_block_size = 1;
+  constexpr int solver = 1;
+  constexpr int max_steps_line_search = 0;
+  auto [mean_est, chol_est] = laplace_latent_tol_solve(
+      poisson_log_likelihood{}, std::forward_as_tuple(y), hessian_block_size,
+      stan::math::test::diagonal_kernel_functor{},
+      std::forward_as_tuple(phi(0), phi(1)),
+      std::make_tuple(theta_0, tolerance, max_num_steps, solver,
+                      max_steps_line_search, true),
+      nullptr);
+  constexpr double tol = 1e-6;
+  EXPECT_EQ(2, mean_est.size());
+  EXPECT_NEAR(theta_root(0), mean_est(0), tol);
+  EXPECT_NEAR(theta_root(1), mean_est(1), tol);
+  EXPECT_NEAR(0.0, chol_est(0, 1), 1e-12);  // check lower triangular matrix
+  Eigen::MatrixXd Sigma_est = chol_est * chol_est.transpose();
+  EXPECT_NEAR(K_laplace(0, 0), Sigma_est(0, 0), tol);
+  EXPECT_NEAR(K_laplace(1, 1), Sigma_est(1, 1), tol);
+  EXPECT_NEAR(K_laplace(0, 1), Sigma_est(0, 1), tol);
+  EXPECT_NEAR(K_laplace(1, 0), Sigma_est(1, 0), tol);
+}
+
+TEST_F(laplace_count_two_dim_diag_test,
+       latent_solve_singular_covariance_throws) {
+  using stan::math::laplace_latent_solve;
+  EXPECT_THROW(({
+                 laplace_latent_solve(
+                     poisson_log_likelihood{}, std::forward_as_tuple(y), 1,
+                     stan::math::test::diagonal_kernel_functor{},
+                     std::forward_as_tuple(0.0, phi(1)),  // singular covariance
+                     nullptr);
+               }),
+               std::domain_error);
+}