Refactor geometric to delegate to neg_binomial(1, beta)

Refactor lpmf, cdf, lcdf, lccdf, and rng to delegate to the
corresponding neg_binomial functions with alpha=1 and
beta=theta/(1-theta), as suggested in review.

Handle theta=1 (where beta diverges) with early returns.
Add autodiff test HPPs for lpmf, lcdf, and lccdf.

Author: yasumorishima

stan/math/prim/prob/geometric_cdf.hpp

@@ -3,40 +3,37 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/constants.hpp>
-#include <stan/math/prim/fun/exp.hpp>
-#include <stan/math/prim/fun/log1m.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/fun/value_of.hpp>
-#include <stan/math/prim/functor/partials_propagator.hpp>
-#include <limits>
+#include <stan/math/prim/prob/neg_binomial_cdf.hpp>
+#include <stan/math/prim/fun/elt_divide.hpp>
+#include <stan/math/prim/fun/subtract.hpp>
 
 namespace stan {
 namespace math {
 
 /** \ingroup prob_dists
- * Returns the CDF of the geometric distribution with success probability
- * parameter. Given containers of matching sizes, returns the product of
- * probabilities.
+ * Returns the CDF of the geometric distribution. Given containers of
+ * matching sizes, returns the product of probabilities.
  *
- * CDF: F(n | theta) = 1 - (1 - theta)^(n + 1)
+ * Delegates to the negative binomial CDF with alpha = 1 and
+ * beta = theta / (1 - theta).
  *
  * @tparam T_n type of outcome variable
  * @tparam T_prob type of success probability parameter
  *
  * @param n outcome variable (number of failures before first success)
  * @param theta success probability parameter
  * @return probability or product of probabilities
- * @throw std::domain_error if theta is not in [0, 1]
+ * @throw std::domain_error if theta is not in (0, 1]
  * @throw std::invalid_argument if container sizes mismatch
  */
 template <typename T_n, typename T_prob>
 inline return_type_t<T_prob> geometric_cdf(const T_n& n,
                                            const T_prob& theta) {
-  using T_partials_return = partials_return_t<T_n, T_prob>;
   using T_n_ref = ref_type_t<T_n>;
   using T_prob_ref = ref_type_t<T_prob>;
   static constexpr const char* function = "geometric_cdf";
@@ -53,42 +50,39 @@ inline return_type_t<T_prob> geometric_cdf(const T_n& n,
                 value_of(theta_ref), 0.0, 1.0);
 
   scalar_seq_view<T_n_ref> n_vec(n_ref);
-  scalar_seq_view<T_prob_ref> theta_vec(theta_ref);
-  const size_t max_size_seq_view = max_size(n_ref, theta_ref);
-
   for (int i = 0; i < stan::math::size(n); i++) {
     if (n_vec.val(i) < 0) {
       return 0.0;
     }
   }
 
-  T_partials_return cdf(1.0);
-  auto ops_partials = make_partials_propagator(theta_ref);
-  for (size_t i = 0; i < max_size_seq_view; i++) {
-    const auto theta_val = theta_vec.val(i);
-    const auto n_val = n_vec.val(i);
-    const auto np1 = n_val + 1.0;
-    const auto log1m_theta = log1m(theta_val);
-    const auto ccdf_i = stan::math::exp(np1 * log1m_theta);
-    const auto cdf_i = 1.0 - ccdf_i;
-
-    cdf *= cdf_i;
-
-    if constexpr (is_autodiff_v<T_prob>) {
-      if (cdf_i > 0.0 && theta_val > 0.0) {
-        partials<0>(ops_partials)[i]
-            += np1 * ccdf_i / ((1.0 - theta_val) * cdf_i);
-      }
+  // theta = 1 => CDF is always 1 for n >= 0
+  scalar_seq_view<T_prob_ref> theta_vec(theta_ref);
+  bool all_theta_one = true;
+  for (size_t i = 0; i < stan::math::size(theta); i++) {
+    if (value_of(theta_vec[i]) != 1.0) {
+      all_theta_one = false;
+      break;
     }
   }
+  if (all_theta_one) {
+    return 1.0;
+  }
 
-  if constexpr (is_autodiff_v<T_prob>) {
-    for (size_t i = 0; i < stan::math::size(theta); ++i) {
-      partials<0>(ops_partials)[i] *= cdf;
+  if constexpr (is_stan_scalar_v<T_prob>) {
+    const auto beta = theta_ref / (1.0 - theta_ref);
+    return neg_binomial_cdf(n_ref, 1, beta);
+  } else if constexpr (is_std_vector_v<T_prob>) {
+    std::vector<value_type_t<T_prob>> beta;
+    beta.reserve(stan::math::size(theta));
+    for (size_t i = 0; i < stan::math::size(theta); i++) {
+      beta.push_back(theta_vec[i] / (1.0 - theta_vec[i]));
     }
+    return neg_binomial_cdf(n_ref, 1, beta);
+  } else {
+    const auto beta = elt_divide(theta_ref, subtract(1.0, theta_ref));
+    return neg_binomial_cdf(n_ref, 1, beta);
   }
-
-  return ops_partials.build(cdf);
 }
 
 }  // namespace math

stan/math/prim/prob/geometric_lccdf.hpp

@@ -3,38 +3,37 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/constants.hpp>
-#include <stan/math/prim/fun/log1m.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/fun/value_of.hpp>
-#include <stan/math/prim/functor/partials_propagator.hpp>
+#include <stan/math/prim/prob/neg_binomial_lccdf.hpp>
+#include <stan/math/prim/fun/elt_divide.hpp>
+#include <stan/math/prim/fun/subtract.hpp>
 
 namespace stan {
 namespace math {
 
 /** \ingroup prob_dists
- * Returns the log CCDF of the geometric distribution with success probability
- * parameter. Given containers of matching sizes, returns the log sum of
- * probabilities.
+ * Returns the log CCDF of the geometric distribution. Given containers of
+ * matching sizes, returns the log sum of probabilities.
  *
- * log CCDF: (n + 1) * log(1 - theta)
+ * Delegates to the negative binomial log CCDF with alpha = 1 and
+ * beta = theta / (1 - theta).
  *
  * @tparam T_n type of outcome variable
  * @tparam T_prob type of success probability parameter
  *
  * @param n outcome variable (number of failures before first success)
  * @param theta success probability parameter
  * @return log complementary probability or log sum
- * @throw std::domain_error if theta is not in [0, 1]
+ * @throw std::domain_error if theta is not in (0, 1]
  * @throw std::invalid_argument if container sizes mismatch
  */
 template <typename T_n, typename T_prob>
 inline return_type_t<T_prob> geometric_lccdf(const T_n& n,
                                              const T_prob& theta) {
-  using T_partials_return = partials_return_t<T_n, T_prob>;
   using T_n_ref = ref_type_t<T_n>;
   using T_prob_ref = ref_type_t<T_prob>;
   static constexpr const char* function = "geometric_lccdf";
@@ -51,31 +50,35 @@ inline return_type_t<T_prob> geometric_lccdf(const T_n& n,
                 value_of(theta_ref), 0.0, 1.0);
 
   scalar_seq_view<T_n_ref> n_vec(n_ref);
-  scalar_seq_view<T_prob_ref> theta_vec(theta_ref);
-  const size_t max_size_seq_view = max_size(n_ref, theta_ref);
-
   for (int i = 0; i < stan::math::size(n); i++) {
     if (n_vec.val(i) < 0) {
       return 0.0;
     }
   }
 
-  T_partials_return log_ccdf(0.0);
-  auto ops_partials = make_partials_propagator(theta_ref);
-  for (size_t i = 0; i < max_size_seq_view; i++) {
-    const auto theta_val = theta_vec.val(i);
-    const auto n_val = n_vec.val(i);
-    const auto np1 = n_val + 1.0;
-    const auto log1m_theta = log1m(theta_val);
-
-    log_ccdf += np1 * log1m_theta;
-
-    if constexpr (is_autodiff_v<T_prob>) {
-      partials<0>(ops_partials)[i] += -np1 / (1.0 - theta_val);
+  // theta = 1 => CCDF = 0 for n >= 0, log CCDF = -inf
+  scalar_seq_view<T_prob_ref> theta_vec(theta_ref);
+  const size_t max_sz = max_size(n_ref, theta_ref);
+  for (size_t i = 0; i < max_sz; i++) {
+    if (value_of(theta_vec[i]) == 1.0 && n_vec.val(i) >= 0) {
+      return negative_infinity();
     }
   }
 
-  return ops_partials.build(log_ccdf);
+  if constexpr (is_stan_scalar_v<T_prob>) {
+    const auto beta = theta_ref / (1.0 - theta_ref);
+    return neg_binomial_lccdf(n_ref, 1, beta);
+  } else if constexpr (is_std_vector_v<T_prob>) {
+    std::vector<value_type_t<T_prob>> beta;
+    beta.reserve(stan::math::size(theta));
+    for (size_t i = 0; i < stan::math::size(theta); i++) {
+      beta.push_back(theta_vec[i] / (1.0 - theta_vec[i]));
+    }
+    return neg_binomial_lccdf(n_ref, 1, beta);
+  } else {
+    const auto beta = elt_divide(theta_ref, subtract(1.0, theta_ref));
+    return neg_binomial_lccdf(n_ref, 1, beta);
+  }
 }
 
 }  // namespace math

stan/math/prim/prob/geometric_lcdf.hpp

@@ -3,40 +3,37 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/constants.hpp>
-#include <stan/math/prim/fun/exp.hpp>
-#include <stan/math/prim/fun/log1m.hpp>
-#include <stan/math/prim/fun/log1m_exp.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/fun/value_of.hpp>
-#include <stan/math/prim/functor/partials_propagator.hpp>
+#include <stan/math/prim/prob/neg_binomial_lcdf.hpp>
+#include <stan/math/prim/fun/elt_divide.hpp>
+#include <stan/math/prim/fun/subtract.hpp>
 
 namespace stan {
 namespace math {
 
 /** \ingroup prob_dists
- * Returns the log CDF of the geometric distribution with success probability
- * parameter. Given containers of matching sizes, returns the log sum of
- * probabilities.
+ * Returns the log CDF of the geometric distribution. Given containers of
+ * matching sizes, returns the log sum of probabilities.
  *
- * log CDF: log(1 - (1 - theta)^(n + 1))
+ * Delegates to the negative binomial log CDF with alpha = 1 and
+ * beta = theta / (1 - theta).
  *
  * @tparam T_n type of outcome variable
  * @tparam T_prob type of success probability parameter
  *
  * @param n outcome variable (number of failures before first success)
  * @param theta success probability parameter
  * @return log probability or log sum of probabilities
- * @throw std::domain_error if theta is not in [0, 1]
+ * @throw std::domain_error if theta is not in (0, 1]
  * @throw std::invalid_argument if container sizes mismatch
  */
 template <typename T_n, typename T_prob>
 inline return_type_t<T_prob> geometric_lcdf(const T_n& n,
                                             const T_prob& theta) {
-  using T_partials_return = partials_return_t<T_n, T_prob>;
   using T_n_ref = ref_type_t<T_n>;
   using T_prob_ref = ref_type_t<T_prob>;
   static constexpr const char* function = "geometric_lcdf";
@@ -53,36 +50,39 @@ inline return_type_t<T_prob> geometric_lcdf(const T_n& n,
                 value_of(theta_ref), 0.0, 1.0);
 
   scalar_seq_view<T_n_ref> n_vec(n_ref);
-  scalar_seq_view<T_prob_ref> theta_vec(theta_ref);
-  const size_t max_size_seq_view = max_size(n_ref, theta_ref);
-
   for (int i = 0; i < stan::math::size(n); i++) {
     if (n_vec.val(i) < 0) {
       return negative_infinity();
     }
   }
 
-  T_partials_return log_cdf(0.0);
-  auto ops_partials = make_partials_propagator(theta_ref);
-  for (size_t i = 0; i < max_size_seq_view; i++) {
-    const auto theta_val = theta_vec.val(i);
-    const auto n_val = n_vec.val(i);
-    const auto np1 = n_val + 1.0;
-    const auto log1m_theta = log1m(theta_val);
-    const auto log_ccdf = np1 * log1m_theta;
-
-    log_cdf += log1m_exp(log_ccdf);
-
-    if constexpr (is_autodiff_v<T_prob>) {
-      const auto ccdf = stan::math::exp(log_ccdf);
-      if (theta_val > 0.0 && ccdf < 1.0) {
-        partials<0>(ops_partials)[i]
-            += np1 * ccdf / ((1.0 - theta_val) * (1.0 - ccdf));
-      }
+  // theta = 1 => log CDF is 0 for n >= 0
+  scalar_seq_view<T_prob_ref> theta_vec(theta_ref);
+  bool all_theta_one = true;
+  for (size_t i = 0; i < stan::math::size(theta); i++) {
+    if (value_of(theta_vec[i]) != 1.0) {
+      all_theta_one = false;
+      break;
     }
   }
+  if (all_theta_one) {
+    return 0.0;
+  }
 
-  return ops_partials.build(log_cdf);
+  if constexpr (is_stan_scalar_v<T_prob>) {
+    const auto beta = theta_ref / (1.0 - theta_ref);
+    return neg_binomial_lcdf(n_ref, 1, beta);
+  } else if constexpr (is_std_vector_v<T_prob>) {
+    std::vector<value_type_t<T_prob>> beta;
+    beta.reserve(stan::math::size(theta));
+    for (size_t i = 0; i < stan::math::size(theta); i++) {
+      beta.push_back(theta_vec[i] / (1.0 - theta_vec[i]));
+    }
+    return neg_binomial_lcdf(n_ref, 1, beta);
+  } else {
+    const auto beta = elt_divide(theta_ref, subtract(1.0, theta_ref));
+    return neg_binomial_lcdf(n_ref, 1, beta);
+  }
 }
 
 }  // namespace math