Merge pull request #3178 from stan-dev/pois-binom-sigs

Fix blocking issues for exposing poisson-binomial in Stan

Author: andrjohns

stan/math/prim/fun/poisson_binomial_log_probs.hpp

@@ -5,6 +5,7 @@
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/log1m.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/log_mix.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
 #include <stan/math/prim/fun/vector_seq_view.hpp>
 
@@ -22,34 +23,34 @@ 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);
 
-  Eigen::Matrix<T_scalar, Eigen::Dynamic, Eigen::Dynamic> alpha(size_theta + 1,
-                                                                y + 1);
+  Eigen::Matrix<T_scalar, Eigen::Dynamic, Eigen::Dynamic> alpha(y + 1,
+                                                                size_theta + 1);
 
   // alpha[i, j] = log prob of j successes in first i trials
   alpha(0, 0) = 0.0;
   for (int i = 0; i < size_theta; ++i) {
     // no success in i trials
-    alpha(i + 1, 0) = alpha(i, 0) + log1m_theta[i];
+    alpha(0, i + 1) = alpha(0, i) + log1m_theta[i];
 
     // 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(j + 1, i + 1) = log_mix(theta[i], alpha(j, i), alpha(j + 1, i));
     }
 
     // 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);
+  return alpha.col(size_theta);
 }
 
 template <typename T_y, typename T_theta, require_vt_integral<T_y>* = nullptr>

stan/math/prim/fun/size_mvt.hpp

@@ -26,6 +26,11 @@ int64_t size_mvt(const ScalarT& /* unused */) {
   throw std::invalid_argument("size_mvt passed to an unrecognized type.");
 }
 
+template <typename ScalarT, require_stan_scalar_t<ScalarT>* = nullptr>
+int64_t size_mvt(const std::vector<ScalarT>& /* unused */) {
+  return 1;
+}
+
 template <typename MatrixT, require_matrix_t<MatrixT>* = nullptr>
 int64_t size_mvt(const MatrixT& /* unused */) {
   return 1;
@@ -36,6 +41,11 @@ int64_t size_mvt(const std::vector<MatrixT>& x) {
   return x.size();
 }
 
+template <typename StdVectorT, require_std_vector_t<StdVectorT>* = nullptr>
+int64_t size_mvt(const std::vector<StdVectorT>& x) {
+  return x.size();
+}
+
 }  // namespace math
 }  // namespace stan
 #endif

stan/math/prim/fun/vector_seq_view.hpp

@@ -6,6 +6,16 @@
 #include <vector>
 
 namespace stan {
+namespace internal {
+template <typename T>
+using is_matrix_or_std_vector
+    = math::disjunction<is_matrix<T>, is_std_vector<T>>;
+
+template <typename T>
+using is_scalar_container = math::disjunction<
+    is_matrix<T>,
+    math::conjunction<is_std_vector<T>, is_stan_scalar<value_type_t<T>>>>;
+}  // namespace internal
 
 /**
  * This class provides a low-cost wrapper for situations where you either need
@@ -33,7 +43,7 @@ class vector_seq_view;
  * @tparam T the type of the underlying Vector
  */
 template <typename T>
-class vector_seq_view<T, require_matrix_t<T>> {
+class vector_seq_view<T, require_t<internal::is_scalar_container<T>>> {
  public:
   explicit vector_seq_view(const T& m) : m_(m) {}
   static constexpr auto size() { return 1; }
@@ -46,19 +56,13 @@ class vector_seq_view<T, require_matrix_t<T>> {
 
   template <typename C = T, require_st_autodiff<C>* = nullptr>
   inline auto val(size_t /* i */) const noexcept {
-    return m_.val();
+    return value_of(m_);
   }
 
  private:
   const ref_type_t<T> m_;
 };
 
-namespace internal {
-template <typename T>
-using is_matrix_or_std_vector
-    = math::disjunction<is_matrix<T>, is_std_vector<T>>;
-}
-
 /**
  * This class provides a low-cost wrapper for situations where you either need
  * an Eigen Vector or RowVector or a std::vector of them and you want to be

stan/math/prim/prob/poisson_binomial_rng.hpp

@@ -21,17 +21,17 @@ namespace math {
  * @return a Poisson binomial distribution random variable
  * @throw std::domain_error if theta is not a valid probability
  */
-template <typename T_theta, typename RNG,
-          require_eigen_vt<std::is_arithmetic, T_theta>* = nullptr>
+template <typename T_theta, typename RNG>
 inline int poisson_binomial_rng(const T_theta& theta, RNG& rng) {
   static constexpr const char* function = "poisson_binomial_rng";
-  check_finite(function, "Probability parameters", theta);
-  check_bounded(function, "Probability parameters", value_of(theta), 0.0, 1.0);
+  ref_type_t<T_theta> theta_ref = theta;
+  check_finite(function, "Probability parameters", theta_ref);
+  check_bounded(function, "Probability parameters", theta_ref, 0.0, 1.0);
 
   int y = 0;
   for (size_t i = 0; i < theta.size(); ++i) {
     boost::variate_generator<RNG&, boost::bernoulli_distribution<> >
-        bernoulli_rng(rng, boost::bernoulli_distribution<>(theta(i)));
+        bernoulli_rng(rng, boost::bernoulli_distribution<>(theta_ref[i]));
     y += bernoulli_rng();
   }
 

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

@@ -6,24 +6,41 @@
 #include <vector>
 
 TEST(ProbDistributionsPoissonBinomial, lpmf_length_0_length_1) {
+  using stan::math::poisson_binomial_lpmf;
+  using stan::math::to_array_1d;
+  using stan::math::to_row_vector;
+
   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_array_1d;
+  using stan::math::to_row_vector;
+
   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 +58,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_array_1d;
+  using stan::math::to_row_vector;
   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_array_1d;
+  using stan::math::to_row_vector;
   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,