laplace_base_rng.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_MIX_FUNCTOR_LAPLACE_BASE_RNG_HPP
#define STAN_MATH_MIX_FUNCTOR_LAPLACE_BASE_RNG_HPP
 
#include <stan/math/prim/fun/Eigen.hpp>
#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 {
 
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 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);
      },
      std::forward<CovarArgs>(covar_args));
  auto md_est = internal::laplace_marginal_density_est(
      ll_fun, std::forward<LLArgs>(ll_args), covariance_train, options, msgs);
  Eigen::VectorXd mean_train = covariance_train * md_est.theta_grad;
  Eigen::MatrixXd Sigma;
  if (md_est.solver_used == 1) {
    Eigen::MatrixXd V_dec
        = md_est.L.template triangularView<Eigen::Lower>().solve(
            md_est.W_r * covariance_train);
    Sigma = covariance_train - V_dec.transpose() * V_dec;
  } else if (md_est.solver_used == 2) {
    Eigen::MatrixXd C = md_est.L.template triangularView<Eigen::Lower>().solve(
        md_est.K_root.transpose());
    Sigma = C.transpose() * C;
  } else {
    Sigma = covariance_train
            - covariance_train
                  * (md_est.W_r
                     - md_est.W_r
                           * md_est.LU.solve(covariance_train * md_est.W_r))
                  * covariance_train;
  }
  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);
  }
}
 
}  // namespace math
}  // namespace stan
 
#endif