Merge pull request #3355 from stan-dev/fix/3353-laplace_bernouilli_logit

update signature for laplace bernoulli logit functions.

Author: SteveBronder

stan/math/mix/prob/laplace_latent_bernoulli_logit_rng.hpp

@@ -20,8 +20,8 @@ namespace math {
  * @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] y binary observations.
+ * @param[in] y_index group to which each observation belongs.
  * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * @param[in] hessian_block_size Block size for the Hessian approximation with
@@ -33,15 +33,15 @@ namespace math {
 template <typename Mean, typename CovarFun, typename CovarArgs,
           typename OpsTuple, typename RNG>
 inline Eigen::VectorXd laplace_latent_tol_bernoulli_logit_rng(
-    const std::vector<int>& y, const std::vector<int>& n_samples, Mean&& mean,
+    const std::vector<int>& y, const std::vector<int>& y_index, Mean&& mean,
     int hessian_block_size, CovarFun&& covariance_function,
     CovarArgs&& covar_args, OpsTuple&& ops, RNG& rng, std::ostream* msgs) {
   auto options
       = internal::tuple_to_laplace_options(std::forward<OpsTuple>(ops));
   options.hessian_block_size = hessian_block_size;
   return laplace_base_rng(
       bernoulli_logit_likelihood{},
-      std::forward_as_tuple(to_vector(y), n_samples, std::forward<Mean>(mean)),
+      std::forward_as_tuple(to_vector(y), y_index, std::forward<Mean>(mean)),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), std::move(options), rng, msgs);
 }
@@ -58,8 +58,8 @@ inline Eigen::VectorXd laplace_latent_tol_bernoulli_logit_rng(
  * @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] y binary observations
+ * @param[in] y_index group to which each observation belongs.
  * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * @param[in] hessian_block_size Block size for the Hessian approximation with
@@ -69,13 +69,13 @@ inline Eigen::VectorXd laplace_latent_tol_bernoulli_logit_rng(
  */
 template <typename Mean, typename CovarFun, typename CovarArgs, typename RNG>
 inline Eigen::VectorXd laplace_latent_bernoulli_logit_rng(
-    const std::vector<int>& y, const std::vector<int>& n_samples, Mean&& mean,
+    const std::vector<int>& y, const std::vector<int>& y_index, Mean&& mean,
     int hessian_block_size, CovarFun&& covariance_function,
     CovarArgs&& covar_args, RNG& rng, std::ostream* msgs) {
   auto options = laplace_options_default{hessian_block_size};
   return laplace_base_rng(
       bernoulli_logit_likelihood{},
-      std::forward_as_tuple(to_vector(y), n_samples, std::forward<Mean>(mean)),
+      std::forward_as_tuple(to_vector(y), y_index, std::forward<Mean>(mean)),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), options, rng, msgs);
 }

stan/math/mix/prob/laplace_marginal_bernoulli_logit_lpmf.hpp

@@ -25,13 +25,36 @@ namespace stan {
 namespace math {
 
 struct bernoulli_logit_likelihood {
+  /**
+   * Returns the lpmf for a Bernoulli with a logit link across
+   * multiple groups. No need to compute the log normalizing constant.
+   * @tparam Theta A type inheriting from `Eigen::EigenBase` with dynamic
+   * sized rows and 1 column.
+   * @tparam YVec A vector type containing integers.
+   * @tparam Mean type of the mean of the latent normal distribution
+   * @param[in] theta log Poisson rate for each group.
+   * @param[in] y binary observations
+   * @param[in] y_index group to which each observation belongs
+   * return lpmf for a Poisson with a log link.
+   * @param[in] mean the mean of the latent normal variable
+   * \msg_arg
+   */
   template <typename ThetaVec, typename YVec, typename Mean>
   inline auto operator()(const ThetaVec& theta, const YVec& y,
-                         const std::vector<int>& delta_int, Mean&& mean,
-                         std::ostream* pstream) const {
+                         const std::vector<int>& y_index, Mean&& mean,
+                         std::ostream* msgs) const {
+    Eigen::VectorXd counts_per_group = Eigen::VectorXd::Zero(theta.size());
+    Eigen::VectorXd n_per_group = Eigen::VectorXd::Zero(theta.size());
+
+    for (int i = 0; i < theta.size(); i++) {
+      counts_per_group(y_index[i] - 1) += y[i];
+      n_per_group(y_index[i] - 1) += 1;
+    }
+
     auto theta_offset = to_ref(add(theta, mean));
-    return sum(elt_multiply(theta_offset, y)
-               - elt_multiply(to_vector(delta_int), log1p_exp(theta_offset)));
+
+    return sum(elt_multiply(theta_offset, counts_per_group)
+               - elt_multiply(to_vector(n_per_group), log1p_exp(theta_offset)));
   }
 };
 
@@ -47,9 +70,8 @@ struct bernoulli_logit_likelihood {
  * 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.
+ * @param[in] y binary observations
+ * @param[in] y_index group to which each observation belongs
  * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * @param[in] hessian_block_size Block size for the Hessian approximation with
@@ -60,15 +82,15 @@ struct bernoulli_logit_likelihood {
 template <bool propto = false, typename Mean, typename CovarFun,
           typename CovarArgs, typename OpsTuple>
 inline auto laplace_marginal_tol_bernoulli_logit_lpmf(
-    const std::vector<int>& y, const std::vector<int>& n_samples, Mean&& mean,
+    const std::vector<int>& y, const std::vector<int>& y_index, Mean&& mean,
     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;
   return laplace_marginal_density(
       bernoulli_logit_likelihood{},
-      std::forward_as_tuple(to_vector(y), n_samples, std::forward<Mean>(mean)),
+      std::forward_as_tuple(to_vector(y), y_index, std::forward<Mean>(mean)),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), std::move(options), msgs);
 }
@@ -83,9 +105,8 @@ inline auto laplace_marginal_tol_bernoulli_logit_lpmf(
  * @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.
+ * @param[in] y binary observations
+ * @param[in] y_index group to which each observation belongs
  * @param[in] mean the mean of the latent normal variable.
  * \laplace_common_args
  * @param[in] hessian_block_size Block size for the Hessian approximation with
@@ -95,13 +116,13 @@ inline auto laplace_marginal_tol_bernoulli_logit_lpmf(
 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, Mean&& mean,
+    const std::vector<int>& y, const std::vector<int>& y_index, Mean&& mean,
     int hessian_block_size, CovarFun&& covariance_function,
     CovarArgs&& covar_args, std::ostream* msgs) {
   auto options = laplace_options_default{hessian_block_size};
   return laplace_marginal_density(
       bernoulli_logit_likelihood{},
-      std::forward_as_tuple(to_vector(y), n_samples, std::forward<Mean>(mean)),
+      std::forward_as_tuple(to_vector(y), y_index, std::forward<Mean>(mean)),
       std::forward<CovarFun>(covariance_function),
       std::forward<CovarArgs>(covar_args), options, msgs);
 }

test/unit/math/laplace/laplace_bernoulli_logit_rng_test.cpp

@@ -66,8 +66,9 @@ TEST(laplace_bernoulli_logit_rng, two_dim_diag) {
 
   Eigen::VectorXd theta_0{{0, 0}};
   Eigen::VectorXd phi{{3, 2}};
-  std::vector<int> n_samples = {1, 1};
-  std::vector<int> sums = {1, 0};
+  std::vector<int> y_index = {1, 2};
+  // std::vector<int> n_samples = {1, 1};
+  std::vector<int> y = {1, 0};
   Eigen::VectorXd ye{{1, 1}};
   Eigen::VectorXd mean{{0, 0}};
   std::vector<double> d0;
@@ -76,7 +77,7 @@ TEST(laplace_bernoulli_logit_rng, two_dim_diag) {
   boost::random::mt19937 rng;
   rng.seed(1954);
   Eigen::MatrixXd theta_pred = laplace_latent_bernoulli_logit_rng(
-      sums, n_samples, mean, 1, diagonal_kernel_functor{},
+      y, y_index, mean, 1, diagonal_kernel_functor{},
       std::forward_as_tuple(phi(0), phi(1)), rng, nullptr);
 
   // Compute exact mean and covariance

test/unit/math/laplace/laplace_marginal_bernoulli_logit_lpmf_test.cpp

@@ -36,15 +36,17 @@ TEST_P(laplace_marginal_bernoulli_logit_lpmf, phi_dim500) {
   for (int i = 0; i < dim_theta; i++) {
     x[i] = Eigen::VectorXd{{x1[i], x2[i]}};
   }
-  std::vector<int> n_samples = stan::math::rep_array(1, dim_theta);
+  std::vector<int> y_index;
+  y_index.reserve(dim_theta);
+  for (int i = 1; i <= dim_theta; i++) {
+    y_index.push_back(i);
+  }
   Eigen::VectorXd theta_0 = Eigen::VectorXd::Zero(dim_theta);
   Eigen::VectorXd mean = Eigen::VectorXd::Zero(dim_theta);
-  std::vector<double> delta;
-  std::vector<int> delta_int;
   Eigen::Matrix<double, Eigen::Dynamic, 1> phi_dbl{{1.6, 1}};
   using stan::math::test::sqr_exp_kernel_functor;
   double target = laplace_marginal_bernoulli_logit_lpmf(
-      y, n_samples, 0, hessian_block_size, sqr_exp_kernel_functor{},
+      y, y_index, 0, hessian_block_size, sqr_exp_kernel_functor{},
       std::forward_as_tuple(x, phi_dbl(0), phi_dbl(1)), nullptr);
   // Benchmark against gpstuff.
   constexpr double tol = 8e-4;
@@ -56,7 +58,7 @@ TEST_P(laplace_marginal_bernoulli_logit_lpmf, phi_dim500) {
   auto f = [&](auto&& alpha, auto&& rho) {
     try {
       return laplace_marginal_tol_bernoulli_logit_lpmf(
-          y, n_samples, mean, hessian_block_size, sqr_exp_kernel_functor{},
+          y, y_index, mean, hessian_block_size, sqr_exp_kernel_functor{},
           std::forward_as_tuple(x, alpha, rho),
           std::make_tuple(theta_0, tolerance, max_num_steps, solver_num,
                           max_steps_line_search, true),