1 #ifndef STAN_MATH_PRIM_FUN_POISSON_BINOMIAL_LOG_PROBS_HPP
2 #define STAN_MATH_PRIM_FUN_POISSON_BINOMIAL_LOG_PROBS_HPP
3
4 #include <stan/math/prim/meta.hpp>
5 #include <stan/math/prim/fun/log.hpp>
6 #include <stan/math/prim/fun/log1m.hpp>
7 #include <stan/math/prim/fun/log_sum_exp.hpp>
8 #include <stan/math/prim/fun/max_size.hpp>
9 #include <stan/math/prim/fun/vector_seq_view.hpp>
10
11 namespace stan {
12 namespace math {
13
14 /**
15 * Returns the last row of the log probability matrix of the Poisson-Binomial
16 * distribution given the number of successes and a vector of success
17 * probabilities.
18 *
19 * @tparam T_theta template expression
20 * @param y numbers of successes
21 * @param theta N-dimensional vector of success probabilities for each trial
22 * @return the last row of the computed log probability matrix
23 */
24 template <typename T_theta, typename T_scalar = scalar_type_t<T_theta>,
25 require_eigen_vector_t<T_theta>* = nullptr>
poisson_binomial_log_probs(int y,const T_theta & theta)26 plain_type_t<T_theta> poisson_binomial_log_probs(int y, const T_theta& theta) {
27 int size_theta = theta.size();
28 plain_type_t<T_theta> log_theta = log(theta);
29 plain_type_t<T_theta> log1m_theta = log1m(theta);
30
31 Eigen::Matrix<T_scalar, Eigen::Dynamic, Eigen::Dynamic> alpha(size_theta + 1,
32 y + 1);
33
34 // alpha[i, j] = log prob of j successes in first i trials
35 alpha(0, 0) = 0.0;
36 for (int i = 0; i < size_theta; ++i) {
37 // no success in i trials
38 alpha(i + 1, 0) = alpha(i, 0) + log1m_theta[i];
39
40 // 0 < j < i successes in i trials
41 for (int j = 0; j < std::min(y, i); ++j) {
42 alpha(i + 1, j + 1) = log_sum_exp(alpha(i, j) + log_theta[i],
43 alpha(i, j + 1) + log1m_theta[i]);
44 }
45
46 // i successes in i trials
47 if (i < y) {
48 alpha(i + 1, i + 1) = alpha(i, i) + log_theta(i);
49 }
50 }
51
52 return alpha.row(size_theta);
53 }
54
55 template <typename T_y, typename T_theta, require_vt_integral<T_y>* = nullptr>
poisson_binomial_log_probs(const T_y & y,const T_theta & theta)56 auto poisson_binomial_log_probs(const T_y& y, const T_theta& theta) {
57 using T_scalar = scalar_type_t<T_theta>;
58 size_t max_sizes = std::max(stan::math::size(y), size_mvt(theta));
59 std::vector<Eigen::Matrix<T_scalar, Eigen::Dynamic, 1>> result(max_sizes);
60 scalar_seq_view<T_y> y_vec(y);
61 vector_seq_view<T_theta> theta_vec(theta);
62
63 for (size_t i = 0; i < max_sizes; ++i) {
64 result[i] = poisson_binomial_log_probs(y_vec[i], theta_vec[i]);
65 }
66
67 return result;
68 }
69
70 } // namespace math
71 } // namespace stan
72 #endif
73