Fix support for all vector signatures in poisson-binomial

Author: andrjohns

stan/math/prim/fun/poisson_binomial_log_probs.hpp

@@ -22,8 +22,9 @@ namespace math {
  * @return the last row of the computed log probability matrix
  */
 template <typename T_theta, typename T_scalar = scalar_type_t<T_theta>,
-          require_eigen_vector_t<T_theta>* = nullptr>
-plain_type_t<T_theta> poisson_binomial_log_probs(int y, const T_theta& theta) {
+          require_vector_t<T_theta>* = nullptr>
+Eigen::Matrix<T_scalar, 1, -1>
+  poisson_binomial_log_probs(int y, const T_theta& theta) {
   int size_theta = theta.size();
   plain_type_t<T_theta> log_theta = log(theta);
   plain_type_t<T_theta> log1m_theta = log1m(theta);
@@ -39,34 +40,18 @@ plain_type_t<T_theta> poisson_binomial_log_probs(int y, const T_theta& theta) {
 
     // 0 < j < i successes in i trials
     for (int j = 0; j < std::min(y, i); ++j) {
-      alpha(i + 1, j + 1) = log_sum_exp(alpha(i, j) + log_theta[i],
-                                        alpha(i, j + 1) + log1m_theta[i]);
+      alpha(i + 1, j + 1) = log_mix(theta[i], alpha(i, j), alpha(i, j + 1));
     }
 
     // i successes in i trials
     if (i < y) {
-      alpha(i + 1, i + 1) = alpha(i, i) + log_theta(i);
+      alpha(i + 1, i + 1) = alpha(i, i) + log_theta[i];
     }
   }
 
   return alpha.row(size_theta);
 }
 
-template <typename T_y, typename T_theta, require_vt_integral<T_y>* = nullptr>
-auto poisson_binomial_log_probs(const T_y& y, const T_theta& theta) {
-  using T_scalar = scalar_type_t<T_theta>;
-  size_t max_sizes = std::max(stan::math::size(y), size_mvt(theta));
-  std::vector<Eigen::Matrix<T_scalar, Eigen::Dynamic, 1>> result(max_sizes);
-  scalar_seq_view<T_y> y_vec(y);
-  vector_seq_view<T_theta> theta_vec(theta);
-
-  for (size_t i = 0; i < max_sizes; ++i) {
-    result[i] = poisson_binomial_log_probs(y_vec[i], theta_vec[i]);
-  }
-
-  return result;
-}
-
 }  // namespace math
 }  // namespace stan
 #endif

test/unit/math/prim/prob/poisson_binomial_test.cpp

@@ -6,24 +6,39 @@
 #include <vector>
 
 TEST(ProbDistributionsPoissonBinomial, lpmf_length_0_length_1) {
+  using stan::math::poisson_binomial_lpmf;
+  using stan::math::to_row_vector;
+  using stan::math::to_array_1d;
+
   Eigen::VectorXd v0(0);
   Eigen::VectorXd v1(1);
   v1 << 0.4;
 
-  EXPECT_FLOAT_EQ(stan::math::poisson_binomial_lpmf(0, v0), 0.0);
-  EXPECT_FLOAT_EQ(stan::math::poisson_binomial_lpmf(1, v1), std::log(0.4));
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(0, v0), 0.0);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(0, to_row_vector(v0)), 0.0);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(0, to_array_1d(v0)), 0.0);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(1, v1), std::log(0.4));
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(1, to_row_vector(v1)), std::log(0.4));
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(1, to_array_1d(v1)), std::log(0.4));
 }
-
 TEST(ProbDistributionsPoissonBinomial, lpmf_length_0_length_1_vectorial_y) {
+  using stan::math::poisson_binomial_lpmf;
+  using stan::math::to_row_vector;
+  using stan::math::to_array_1d;
+
   Eigen::VectorXd v0(0);
   Eigen::VectorXd v1(1);
   v1 << 0.4;
 
   std::vector<int> y0{0, 0};
   std::vector<int> y1{1, 1};
 
-  EXPECT_FLOAT_EQ(stan::math::poisson_binomial_lpmf(y0, v0), 0.0);
-  EXPECT_FLOAT_EQ(stan::math::poisson_binomial_lpmf(y1, v1), std::log(0.4) * 2);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(y0, v0), 0.0);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(y0, to_row_vector(v0)), 0.0);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(y0, to_array_1d(v0)), 0.0);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(y1, v1), std::log(0.4) * 2);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(y1, to_row_vector(v1)), std::log(0.4) * 2);
+  EXPECT_FLOAT_EQ(poisson_binomial_lpmf(y1, to_array_1d(v1)), std::log(0.4) * 2);
 }
 
 TEST(ProbDistributionsPoissonBinomial, lpmf_works_on_scalar_arguments) {
@@ -41,25 +56,35 @@ TEST(ProbDistributionsPoissonBinomial, lpmf_works_on_scalar_arguments) {
 
 TEST(ProbDistributionsPoissonBinomial, lpmf_works_on_vectorial_y) {
   using stan::math::poisson_binomial_lpmf;
+  using stan::math::to_row_vector;
+  using stan::math::to_array_1d;
   using vec = Eigen::Matrix<double, Eigen::Dynamic, 1>;
 
   vec p(3, 1);
   p << 0.5, 0.2, 0.7;
   std::vector<int> y{2, 2};
 
   EXPECT_NEAR(-0.967584 * 2, poisson_binomial_lpmf(y, p), 1e-6);
+  EXPECT_NEAR(-0.967584 * 2, poisson_binomial_lpmf(y, to_row_vector(p)), 1e-6);
+  EXPECT_NEAR(-0.967584 * 2, poisson_binomial_lpmf(y, to_array_1d(p)), 1e-6);
 }
 
 TEST(ProbDistributionsPoissonBinomial, lpmf_works_on_vectorial_y_and_theta) {
   using stan::math::poisson_binomial_lpmf;
+  using stan::math::to_row_vector;
+  using stan::math::to_array_1d;
   using vec = Eigen::Matrix<double, Eigen::Dynamic, 1>;
 
   vec p(3, 1);
   p << 0.5, 0.2, 0.7;
   std::vector<int> y{2, 0};
   std::vector<vec> ps{p, p};
+  std::vector<Eigen::RowVectorXd> ps_rv{to_row_vector(p), to_row_vector(p)};
+  std::vector<std::vector<double>> ps_std{to_array_1d(p), to_array_1d(p)};
 
   EXPECT_NEAR(-0.967584 - 2.12026, poisson_binomial_lpmf(y, ps), 1e-5);
+  EXPECT_NEAR(-0.967584 - 2.12026, poisson_binomial_lpmf(y, ps_rv), 1e-5);
+  EXPECT_NEAR(-0.967584 - 2.12026, poisson_binomial_lpmf(y, ps_std), 1e-5);
 }
 
 TEST(ProbDistributionsPoissonBinomial,