Merge pull request #3116 from lingium/feature/issue-3115-beta-neg-bin…

…omial-lcdf

add beta_neg_binomial_lcdf

stan/math/prim/prob.hpp

@@ -26,6 +26,7 @@
 #include <stan/math/prim/prob/beta_lcdf.hpp>
 #include <stan/math/prim/prob/beta_lpdf.hpp>
 #include <stan/math/prim/prob/beta_neg_binomial_lccdf.hpp>
+#include <stan/math/prim/prob/beta_neg_binomial_lcdf.hpp>
 #include <stan/math/prim/prob/beta_neg_binomial_lpmf.hpp>
 #include <stan/math/prim/prob/beta_proportion_ccdf_log.hpp>
 #include <stan/math/prim/prob/beta_proportion_cdf_log.hpp>

stan/math/prim/prob/beta_neg_binomial_lccdf.hpp

@@ -34,15 +34,15 @@ namespace math {
  * @param alpha prior success parameter
  * @param beta prior failure parameter
  * @param precision precision for `grad_F32`, default \f$10^{-8}\f$
- * @param max_steps max iteration allowed for `grad_F32`, default \f$10^{-8}\f$
+ * @param max_steps max iteration allowed for `grad_F32`, default \f$10^{8}\f$
  * @return log probability or log sum of probabilities
  * @throw std::domain_error if r, alpha, or beta fails to be positive
  * @throw std::invalid_argument if container sizes mismatch
  */
 template <typename T_n, typename T_r, typename T_alpha, typename T_beta>
 inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lccdf(
     const T_n& n, const T_r& r, const T_alpha& alpha, const T_beta& beta,
-    const double precision = 1e-8, const int max_steps = 1e6) {
+    const double precision = 1e-8, const int max_steps = 1e8) {
   static constexpr const char* function = "beta_neg_binomial_lccdf";
   check_consistent_sizes(
       function, "Failures variable", n, "Number of successes parameter", r,
@@ -100,7 +100,7 @@ inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lccdf(
         std::initializer_list<b_t>{n_dbl + 2.0, a_plus_r + b_plus_n + 1.0},
         1.0);
     auto C = lgamma(r_plus_n + 1.0) + lbeta(a_plus_r, b_plus_n + 1.0)
-             - lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl) - lgamma(n_dbl + 2);
+             - lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl) - lgamma(n_dbl + 2.0);
     log_ccdf += C + stan::math::log(F);
 
     if constexpr (!is_constant_all<T_r, T_alpha, T_beta>::value) {

stan/math/prim/prob/beta_neg_binomial_lcdf.hpp

@@ -0,0 +1,152 @@
+#ifndef STAN_MATH_PRIM_PROB_BETA_NEG_BINOMIAL_LCDF_HPP
+#define STAN_MATH_PRIM_PROB_BETA_NEG_BINOMIAL_LCDF_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/digamma.hpp>
+#include <stan/math/prim/fun/hypergeometric_3F2.hpp>
+#include <stan/math/prim/fun/grad_F32.hpp>
+#include <stan/math/prim/fun/lbeta.hpp>
+#include <stan/math/prim/fun/lgamma.hpp>
+#include <stan/math/prim/fun/max_size.hpp>
+#include <stan/math/prim/fun/scalar_seq_view.hpp>
+#include <stan/math/prim/fun/size.hpp>
+#include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/functor/partials_propagator.hpp>
+#include <cmath>
+
+namespace stan {
+namespace math {
+
+/** \ingroup prob_dists
+ * Returns the log CDF of the Beta-Negative Binomial distribution with given
+ * number of successes, prior success, and prior failure parameters.
+ * Given containers of matching sizes, returns the log sum of probabilities.
+ *
+ * @tparam T_n type of failure parameter
+ * @tparam T_r type of number of successes parameter
+ * @tparam T_alpha type of prior success parameter
+ * @tparam T_beta type of prior failure parameter
+ *
+ * @param n failure parameter
+ * @param r Number of successes parameter
+ * @param alpha prior success parameter
+ * @param beta prior failure parameter
+ * @param precision precision for `grad_F32`, default \f$10^{-8}\f$
+ * @param max_steps max iteration allowed for `grad_F32`, default \f$10^{8}\f$
+ * @return log probability or log sum of probabilities
+ * @throw std::domain_error if r, alpha, or beta fails to be positive
+ * @throw std::invalid_argument if container sizes mismatch
+ */
+template <typename T_n, typename T_r, typename T_alpha, typename T_beta>
+inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lcdf(
+    const T_n& n, const T_r& r, const T_alpha& alpha, const T_beta& beta,
+    const double precision = 1e-8, const int max_steps = 1e8) {
+  static constexpr const char* function = "beta_neg_binomial_lcdf";
+  check_consistent_sizes(
+      function, "Failures variable", n, "Number of successes parameter", r,
+      "Prior success parameter", alpha, "Prior failure parameter", beta);
+  if (size_zero(n, r, alpha, beta)) {
+    return 0;
+  }
+
+  using T_r_ref = ref_type_t<T_r>;
+  T_r_ref r_ref = r;
+  using T_alpha_ref = ref_type_t<T_alpha>;
+  T_alpha_ref alpha_ref = alpha;
+  using T_beta_ref = ref_type_t<T_beta>;
+  T_beta_ref beta_ref = beta;
+  check_positive_finite(function, "Number of successes parameter", r_ref);
+  check_positive_finite(function, "Prior success parameter", alpha_ref);
+  check_positive_finite(function, "Prior failure parameter", beta_ref);
+
+  scalar_seq_view<T_n> n_vec(n);
+  scalar_seq_view<T_r_ref> r_vec(r_ref);
+  scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
+  scalar_seq_view<T_beta_ref> beta_vec(beta_ref);
+  int size_n = stan::math::size(n);
+  size_t max_size_seq_view = max_size(n, r, alpha, beta);
+
+  // Explicit return for extreme values
+  // The gradients are technically ill-defined, but treated as zero
+  for (int i = 0; i < size_n; i++) {
+    if (n_vec.val(i) < 0) {
+      return negative_infinity();
+    }
+  }
+
+  using T_partials_return = partials_return_t<T_n, T_r, T_alpha, T_beta>;
+  T_partials_return log_cdf(0.0);
+  auto ops_partials = make_partials_propagator(r_ref, alpha_ref, beta_ref);
+  for (size_t i = 0; i < max_size_seq_view; i++) {
+    // Explicit return for extreme values
+    // The gradients are technically ill-defined, but treated as zero
+    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
+      return 0.0;
+    }
+    auto n_dbl = n_vec.val(i);
+    auto r_dbl = r_vec.val(i);
+    auto alpha_dbl = alpha_vec.val(i);
+    auto beta_dbl = beta_vec.val(i);
+    auto b_plus_n = beta_dbl + n_dbl;
+    auto r_plus_n = r_dbl + n_dbl;
+    auto a_plus_r = alpha_dbl + r_dbl;
+    using a_t = return_type_t<decltype(b_plus_n), decltype(r_plus_n)>;
+    using b_t = return_type_t<decltype(n_dbl), decltype(a_plus_r),
+                              decltype(b_plus_n)>;
+    auto F = hypergeometric_3F2(
+        std::initializer_list<a_t>{1.0, b_plus_n + 1.0, r_plus_n + 1.0},
+        std::initializer_list<b_t>{n_dbl + 2.0, a_plus_r + b_plus_n + 1.0},
+        1.0);
+    auto C = lgamma(r_plus_n + 1.0) + lbeta(a_plus_r, b_plus_n + 1.0)
+             - lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl) - lgamma(n_dbl + 2.0);
+    auto ccdf = stan::math::exp(C) * F;
+    log_cdf += log1m(ccdf);
+
+    if constexpr (!is_constant_all<T_r, T_alpha, T_beta>::value) {
+      auto chain_rule_term = -ccdf / (1.0 - ccdf);
+      auto digamma_n_r_alpha_beta = digamma(a_plus_r + b_plus_n + 1.0);
+      T_partials_return dF[6];
+      grad_F32<false, !is_constant<T_beta>::value, !is_constant_all<T_r>::value,
+               false, true, false>(dF, 1.0, b_plus_n + 1.0, r_plus_n + 1.0,
+                                   n_dbl + 2.0, a_plus_r + b_plus_n + 1.0, 1.0,
+                                   precision, max_steps);
+
+      if constexpr (!is_constant<T_r>::value || !is_constant<T_alpha>::value) {
+        auto digamma_r_alpha = digamma(a_plus_r);
+        if constexpr (!is_constant<T_r>::value) {
+          auto partial_lccdf = digamma(r_plus_n + 1.0)
+                               + (digamma_r_alpha - digamma_n_r_alpha_beta)
+                               + (dF[2] + dF[4]) / F - digamma(r_dbl);
+          partials<0>(ops_partials)[i] += partial_lccdf * chain_rule_term;
+        }
+        if constexpr (!is_constant<T_alpha>::value) {
+          auto partial_lccdf = digamma_r_alpha - digamma_n_r_alpha_beta
+                               + dF[4] / F - digamma(alpha_dbl);
+          partials<1>(ops_partials)[i] += partial_lccdf * chain_rule_term;
+        }
+      }
+
+      if constexpr (!is_constant<T_alpha>::value
+                    || !is_constant<T_beta>::value) {
+        auto digamma_alpha_beta = digamma(alpha_dbl + beta_dbl);
+        if constexpr (!is_constant<T_alpha>::value) {
+          partials<1>(ops_partials)[i] += digamma_alpha_beta * chain_rule_term;
+        }
+        if constexpr (!is_constant<T_beta>::value) {
+          auto partial_lccdf = digamma(b_plus_n + 1.0) - digamma_n_r_alpha_beta
+                               + (dF[1] + dF[4]) / F
+                               - (digamma(beta_dbl) - digamma_alpha_beta);
+          partials<2>(ops_partials)[i] += partial_lccdf * chain_rule_term;
+        }
+      }
+    }
+  }
+
+  return ops_partials.build(log_cdf);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

test/prob/beta_neg_binomial/beta_neg_binomial_cdf_log_test.hpp

@@ -0,0 +1,92 @@
+// Arguments: Ints, Doubles, Doubles, Doubles
+#include <stan/math/prim/prob/beta_neg_binomial_lcdf.hpp>
+#include <stan/math/prim/fun/lbeta.hpp>
+#include <stan/math/prim/fun/lgamma.hpp>
+
+using stan::math::var;
+using std::numeric_limits;
+using std::vector;
+
+class AgradCdfLogBetaNegBinomial : public AgradCdfLogTest {
+ public:
+  void valid_values(vector<vector<double>>& parameters,
+                    vector<double>& cdf_log) {
+    vector<double> param(4);
+
+    param[0] = 0;    // n
+    param[1] = 1.0;  // r
+    param[2] = 5.0;  // alpha
+    param[3] = 1.0;  // beta
+    parameters.push_back(param);
+    cdf_log.push_back(std::log(0.833333333333333));  // expected cdf_log
+  }
+
+  void invalid_values(vector<size_t>& index, vector<double>& value) {
+    // n
+
+    // r
+    index.push_back(1U);
+    value.push_back(0.0);
+
+    index.push_back(1U);
+    value.push_back(-1.0);
+
+    index.push_back(1U);
+    value.push_back(std::numeric_limits<double>::infinity());
+
+    // alpha
+    index.push_back(2U);
+    value.push_back(0.0);
+
+    index.push_back(2U);
+    value.push_back(-1.0);
+
+    index.push_back(2U);
+    value.push_back(std::numeric_limits<double>::infinity());
+
+    // beta
+    index.push_back(3U);
+    value.push_back(0.0);
+
+    index.push_back(3U);
+    value.push_back(-1.0);
+
+    index.push_back(3U);
+    value.push_back(std::numeric_limits<double>::infinity());
+  }
+
+  // BOUND INCLUDED IN ORDER FOR TEST TO PASS WITH CURRENT FRAMEWORK
+  bool has_lower_bound() { return false; }
+
+  bool has_upper_bound() { return false; }
+
+  template <typename T_n, typename T_r, typename T_size1, typename T_size2,
+            typename T4, typename T5>
+  stan::return_type_t<T_r, T_size1, T_size2> cdf_log(const T_n& n, const T_r& r,
+                                                     const T_size1& alpha,
+                                                     const T_size2& beta,
+                                                     const T4&, const T5&) {
+    return stan::math::beta_neg_binomial_lcdf(n, r, alpha, beta);
+  }
+
+  template <typename T_n, typename T_r, typename T_size1, typename T_size2,
+            typename T4, typename T5>
+  stan::return_type_t<T_r, T_size1, T_size2> cdf_log_function(
+      const T_n& n, const T_r& r, const T_size1& alpha, const T_size2& beta,
+      const T4&, const T5&) {
+    using stan::math::lbeta;
+    using stan::math::lgamma;
+    using stan::math::log_sum_exp;
+    using std::vector;
+
+    vector<stan::return_type_t<T_r, T_size1, T_size2>> lpmf_values;
+
+    for (int i = 0; i <= n; i++) {
+      auto lpmf = lbeta(i + r, alpha + beta) - lbeta(r, alpha)
+                  + lgamma(i + beta) - lgamma(i + 1) - lgamma(beta);
+      lpmf_values.push_back(lpmf);
+    }
+
+    return log_sum_exp(lpmf_values);
+  }
+};