Add mean to laplace helper distributions

Author: WardBrian

stan/math/mix/prob.hpp

@@ -3,13 +3,11 @@
 
 #include <stan/math/mix/prob/laplace_latent_bernoulli_logit_rng.hpp>
 #include <stan/math/mix/prob/laplace_latent_poisson_log_rng.hpp>
-#include <stan/math/mix/prob/laplace_latent_poisson_log_2_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_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>
-#include <stan/math/mix/prob/laplace_marginal_poisson_log_2_lpmf.hpp>
 #include <stan/math/mix/prob/laplace_marginal_poisson_log_lpmf.hpp>
 
 #endif

stan/math/mix/prob/laplace_latent_bernoulli_logit_rng.hpp

@@ -19,30 +19,34 @@ namespace math {
  * where the likelihood is a Bernoulli with logit link.
  * @tparam ThetaVec A type inheriting from `Eigen::EigenBase`
  * with dynamic sized rows and 1 column.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @tparam RNG A valid boost rng type
  * @param[in] y Vector Vector of total number of trials with a positive outcome.
  * @param[in] n_samples Vector of number of trials.
+ * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * \laplace_options
  * \rng_arg
  * \msg_arg
  */
-template <typename ThetaVec, typename CovarFun, typename CovarArgs,
-          typename RNG, require_eigen_t<ThetaVec>* = nullptr>
+template <typename ThetaVec, typename Mean, typename CovarFun,
+          typename CovarArgs, typename RNG,
+          require_all_eigen_vector_t<ThetaVec>* = nullptr>
 inline Eigen::VectorXd laplace_latent_tol_bernoulli_logit_rng(
-    const std::vector<int>& y, const std::vector<int>& n_samples,
+    const std::vector<int>& y, const std::vector<int>& n_samples, Mean&& mean,
     CovarFun&& covariance_function, CovarArgs&& covar_args, ThetaVec&& theta_0,
     const double tolerance, const int max_num_steps,
     const int hessian_block_size, const int solver,
     const int max_steps_line_search, RNG& rng, std::ostream* msgs) {
   laplace_options_user_supplied ops{hessian_block_size,    solver,
                                     max_steps_line_search, tolerance,
                                     max_num_steps,         value_of(theta_0)};
-  return laplace_base_rng(bernoulli_logit_likelihood{},
-                          std::forward_as_tuple(to_vector(y), n_samples),
-                          std::forward<CovarFun>(covariance_function),
-                          std::forward<CovarArgs>(covar_args), ops, rng, msgs);
+  return laplace_base_rng(
+      bernoulli_logit_likelihood{},
+      std::forward_as_tuple(to_vector(y), n_samples, std::forward<Mean>(mean)),
+      std::forward<CovarFun>(covariance_function),
+      std::forward<CovarArgs>(covar_args), ops, rng, msgs);
 }
 
 /**
@@ -54,24 +58,28 @@ inline Eigen::VectorXd laplace_latent_tol_bernoulli_logit_rng(
  * return a multivariate normal random variate sampled
  * from the gaussian approximation of p(theta | y, phi),
  * where the likelihood is a Bernoulli with logit link.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @tparam RNG A valid boost rng type
  * @param[in] y Vector Vector of total number of trials with a positive outcome.
  * @param[in] n_samples Vector of number of trials.
+ * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * \rng_arg
  * \msg_arg
  */
-template <typename CovarFun, typename CovarArgs, typename RNG>
+template <typename Mean, typename CovarFun, typename CovarArgs, typename RNG,
+          require_eigen_vector_t<Mean>* = nullptr>
 inline Eigen::VectorXd laplace_latent_bernoulli_logit_rng(
-    const std::vector<int>& y, const std::vector<int>& n_samples,
+    const std::vector<int>& y, const std::vector<int>& n_samples, Mean&& mean,
     CovarFun&& covariance_function, CovarArgs&& covar_args, RNG& rng,
     std::ostream* msgs) {
-  return laplace_base_rng(bernoulli_logit_likelihood{},
-                          std::forward_as_tuple(to_vector(y), n_samples),
-                          std::forward<CovarFun>(covariance_function),
-                          std::forward<CovarArgs>(covar_args),
-                          laplace_options_default{}, rng, msgs);
+  return laplace_base_rng(
+      bernoulli_logit_likelihood{},
+      std::forward_as_tuple(to_vector(y), n_samples, std::forward<Mean>(mean)),
+      std::forward<CovarFun>(covariance_function),
+      std::forward<CovarArgs>(covar_args), laplace_options_default{}, rng,
+      msgs);
 }
 
 }  // namespace math

stan/math/mix/prob/laplace_latent_neg_binomial_2_log_rng.hpp

@@ -23,31 +23,34 @@ namespace math {
  * @tparam Eta A type for the overdispersion parameter.
  * @tparam ThetaVec A type inheriting from `Eigen::EigenBase`
  * with dynamic sized rows and 1 column.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @tparam RNG A valid boost rng type
  * @param[in] y Observed counts.
  * @param[in] y_index Index indicating which group each observation belongs to.
  * @param[in] eta Overdisperison parameter.
+ * @param[in] mean The mean of the latent normal variable.
  * \laplace_common_args
  * \laplace_options
  * \rng_arg
  * \msg_arg
  */
-template <typename Eta, typename ThetaVec, typename CovarFun,
+template <typename Eta, typename ThetaVec, typename Mean, typename CovarFun,
           typename CovarArgs, typename RNG,
-          require_eigen_t<ThetaVec>* = nullptr>
+          require_eigen_vector_t<ThetaVec>* = nullptr>
 inline Eigen::VectorXd laplace_latent_tol_neg_binomial_2_log_rng(
     const std::vector<int>& y, const std::vector<int>& y_index, Eta&& eta,
-    CovarFun&& covariance_function, CovarArgs&& covar_args, ThetaVec&& theta_0,
-    const double tolerance, const int max_num_steps,
+    Mean&& mean, CovarFun&& covariance_function, CovarArgs&& covar_args,
+    ThetaVec&& theta_0, const double tolerance, const int max_num_steps,
     const int hessian_block_size, const int solver,
     const int max_steps_line_search, RNG& rng, std::ostream* msgs) {
   laplace_options_user_supplied ops{hessian_block_size,    solver,
                                     max_steps_line_search, tolerance,
                                     max_num_steps,         value_of(theta_0)};
   return laplace_base_rng(
       neg_binomial_2_log_likelihood{},
-      std::forward_as_tuple(std::forward<Eta>(eta), y, y_index),
+      std::forward_as_tuple(std::forward<Eta>(eta), y, y_index,
+                            std::forward<Mean>(mean)),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), ops, rng, msgs);
 }
@@ -65,23 +68,27 @@ inline Eigen::VectorXd laplace_latent_tol_neg_binomial_2_log_rng(
  * parameterization of the Negative Binomial.
  *
  * @tparam Eta A type for the overdispersion parameter.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @tparam RNG A valid boost rng type
  * @param[in] y Observed counts.
  * @param[in] y_index Index indicating which group each observation belongs to.
  * @param[in] eta Overdisperison parameter.
+ * @param[in] mean The mean of the latent normal variable.
  * \laplace_common_args
  * \rng_arg
  * \msg_arg
  */
-template <typename Eta, typename CovarFun, typename CovarArgs, typename RNG>
+template <typename Eta, typename Mean, typename CovarFun, typename CovarArgs,
+          typename RNG>
 inline Eigen::VectorXd laplace_latent_neg_binomial_2_log_rng(
     const std::vector<int>& y, const std::vector<int>& y_index, Eta&& eta,
-    CovarFun&& covariance_function, CovarArgs&& covar_args, RNG& rng,
-    std::ostream* msgs) {
+    Mean&& mean, CovarFun&& covariance_function, CovarArgs&& covar_args,
+    RNG& rng, std::ostream* msgs) {
   return laplace_base_rng(
       neg_binomial_2_log_likelihood{},
-      std::forward_as_tuple(std::forward<Eta>(eta), y, y_index),
+      std::forward_as_tuple(std::forward<Eta>(eta), y, y_index,
+                            std::forward<Mean>(mean)),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), laplace_options_default{}, rng,
       msgs);

stan/math/mix/prob/laplace_latent_poisson_log_2_rng.hpp

@@ -19,28 +19,31 @@ namespace math {
  * In this specialized function, the likelihood p(y|theta) is a
  * @tparam ThetaVec A type inheriting from `Eigen::EigenBase`
  * with dynamic sized rows and 1 column.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @tparam RNG A valid boost rng type
  * @param[in] y Observed counts.
  * @param[in] y_index Index indicating which group each observation belongs to.
+ * @param[in] mean The mean of the latent normal variable.
  * \laplace_common_args
  * \laplace_options
  * \rng_arg
  * \msg_arg
  */
-template <typename ThetaVec, typename CovarFun, typename CovarArgs,
-          typename RNG, require_eigen_t<ThetaVec>* = nullptr>
+template <typename ThetaVec, typename Mean, typename CovarFun,
+          typename CovarArgs, typename RNG,
+          require_eigen_vector_t<ThetaVec>* = nullptr>
 inline Eigen::VectorXd laplace_latent_tol_poisson_log_rng(
     const std::vector<int>& y, const std::vector<int>& y_index,
-    CovarFun&& covariance_function, CovarArgs&& covar_args, ThetaVec&& theta_0,
-    const double tolerance, const int max_num_steps,
+    const Mean& mean, CovarFun&& covariance_function, CovarArgs&& covar_args,
+    ThetaVec&& theta_0, const double tolerance, const int max_num_steps,
     const int hessian_block_size, const int solver,
     const int max_steps_line_search, RNG& rng, std::ostream* msgs) {
   laplace_options_user_supplied ops{hessian_block_size,    solver,
                                     max_steps_line_search, tolerance,
                                     max_num_steps,         value_of(theta_0)};
   return laplace_base_rng(poisson_log_likelihood{},
-                          std::forward_as_tuple(y, y_index),
+                          std::forward_as_tuple(y, y_index, mean),
                           std::forward<CovarFun>(covariance_function),
                           std::forward<CovarArgs>(covar_args), ops, rng, msgs);
 }
@@ -55,21 +58,23 @@ inline Eigen::VectorXd laplace_latent_tol_poisson_log_rng(
  * The Laplace approximation is computed using a Newton solver.
  * In this specialized function, the likelihood p(y|theta) is a
  * Poisson with a log link.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @tparam RNG A valid boost rng type
  * @param[in] y Observed counts.
  * @param[in] y_index Index indicating which group each observation belongs to.
+ * @param[in] mean The mean of the latent normal variable.
  * \laplace_common_args
  * \rng_arg
  * \msg_arg
  */
-template <typename CovarFun, typename CovarArgs, typename RNG>
+template <typename CovarFun, typename CovarArgs, typename RNG, typename Mean>
 inline Eigen::VectorXd laplace_latent_poisson_log_rng(
     const std::vector<int>& y, const std::vector<int>& y_index,
-    CovarFun&& covariance_function, CovarArgs&& covar_args, RNG& rng,
-    std::ostream* msgs) {
+    const Mean& mean, CovarFun&& covariance_function, CovarArgs&& covar_args,
+    RNG& rng, std::ostream* msgs) {
   return laplace_base_rng(poisson_log_likelihood{},
-                          std::forward_as_tuple(y, y_index),
+                          std::forward_as_tuple(y, y_index, mean),
                           std::forward<CovarFun>(covariance_function),
                           std::forward<CovarArgs>(covar_args),
                           laplace_options_default{}, rng, msgs);

stan/math/mix/prob/laplace_latent_poisson_log_rng.hpp

@@ -20,12 +20,14 @@ namespace stan {
 namespace math {
 
 struct bernoulli_logit_likelihood {
-  template <typename T_theta, typename YVec>
-  inline auto operator()(const T_theta& theta, const YVec& y,
-                         const std::vector<int>& delta_int,
+  template <typename ThetaVec, typename YVec, typename Mean>
+  inline auto operator()(const ThetaVec& theta, const YVec& y,
+                         const std::vector<int>& delta_int, const Mean& mean,
                          std::ostream* pstream) const {
-    return sum(elt_multiply(theta, y)
-               - elt_multiply(to_vector(delta_int), log(add(1.0, exp(theta)))));
+    auto theta_offset = to_ref(add(theta, mean));
+    return sum(
+        elt_multiply(theta_offset, y)
+        - elt_multiply(to_vector(delta_int), log(add(1.0, exp(theta_offset)))));
   }
 };
 
@@ -39,19 +41,22 @@ struct bernoulli_logit_likelihood {
  * @tparam propto boolean ignored
  * @tparam ThetaVec A type inheriting from `Eigen::EigenBase`
  * with dynamic sized rows and 1 column.
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @param[in] y total counts per group. Second sufficient statistics.
  * @param[in] n_samples number of samples per group. First sufficient
- *            statistics.
+ * statistics.
+ * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * \laplace_options
  * \msg_arg
  */
-template <bool propto = false, typename ThetaVec, typename CovarFun,
-          typename CovarArgs, require_eigen_t<ThetaVec>* = nullptr>
+template <bool propto = false, typename ThetaVec, typename Mean,
+          typename CovarFun, typename CovarArgs,
+          require_eigen_vector_t<ThetaVec>* = nullptr>
 inline auto laplace_marginal_tol_bernoulli_logit_lpmf(
     const std::vector<int>& y, const std::vector<int>& n_samples,
-    CovarFun&& covariance_function, CovarArgs&& covar_args,
+    const Mean& mean, CovarFun&& covariance_function, CovarArgs&& covar_args,
     const ThetaVec& theta_0, double tolerance, int max_num_steps,
     const int hessian_block_size, const int solver,
     const int max_steps_line_search, std::ostream* msgs) {
@@ -60,7 +65,7 @@ inline auto laplace_marginal_tol_bernoulli_logit_lpmf(
                                     max_num_steps,         value_of(theta_0)};
   return laplace_marginal_density(
       bernoulli_logit_likelihood{},
-      std::forward_as_tuple(to_vector(y), n_samples),
+      std::forward_as_tuple(to_vector(y), n_samples, mean),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), ops, msgs);
 }
@@ -73,21 +78,24 @@ inline auto laplace_marginal_tol_bernoulli_logit_lpmf(
  * for more details.
  *
  * @tparam propto boolean ignored
+ * @tparam Mean type of the mean of the latent normal distribution
  * \laplace_common_template_args
  * @param[in] y total counts per group. Second sufficient statistics.
  * @param[in] n_samples number of samples per group. First sufficient
- *            statistics.
+ * statistics.
+ * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * \msg_arg
  */
-template <bool propto = false, typename CovarFun, typename CovarArgs>
+template <bool propto = false, typename Mean, typename CovarFun,
+          typename CovarArgs>
 inline auto laplace_marginal_bernoulli_logit_lpmf(
     const std::vector<int>& y, const std::vector<int>& n_samples,
-    CovarFun&& covariance_function, CovarArgs&& covar_args,
+    const Mean& mean, CovarFun&& covariance_function, CovarArgs&& covar_args,
     std::ostream* msgs) {
   return laplace_marginal_density(
       bernoulli_logit_likelihood{},
-      std::forward_as_tuple(to_vector(y), n_samples),
+      std::forward_as_tuple(to_vector(y), n_samples, mean),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), laplace_options_default{}, msgs);
 }