Add geometric distribution (lpmf, cdf, lcdf, lccdf, rng)

stan/math/prim/prob.hpp

@@ -113,6 +113,11 @@
 #include <stan/math/prim/prob/gamma_rng.hpp>
 #include <stan/math/prim/prob/gaussian_dlm_obs_lpdf.hpp>
 #include <stan/math/prim/prob/gaussian_dlm_obs_rng.hpp>
+#include <stan/math/prim/prob/geometric_cdf.hpp>
+#include <stan/math/prim/prob/geometric_lccdf.hpp>
+#include <stan/math/prim/prob/geometric_lcdf.hpp>
+#include <stan/math/prim/prob/geometric_lpmf.hpp>
+#include <stan/math/prim/prob/geometric_rng.hpp>
 #include <stan/math/prim/prob/gumbel_ccdf_log.hpp>
 #include <stan/math/prim/prob/gumbel_cdf.hpp>
 #include <stan/math/prim/prob/gumbel_cdf_log.hpp>

stan/math/prim/prob/geometric_cdf.hpp

@@ -0,0 +1,94 @@
+#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_CDF_HPP
+#define STAN_MATH_PRIM_PROB_GEOMETRIC_CDF_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/any.hpp>
+#include <stan/math/prim/fun/as_value_column_array_or_scalar.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/elt_divide.hpp>
+#include <stan/math/prim/fun/exp.hpp>
+#include <stan/math/prim/fun/expm1.hpp>
+#include <stan/math/prim/fun/log1m.hpp>
+#include <stan/math/prim/fun/prod.hpp>
+#include <stan/math/prim/fun/select.hpp>
+#include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/sum.hpp>
+#include <stan/math/prim/fun/value_of.hpp>
+#include <stan/math/prim/functor/partials_propagator.hpp>
+
+namespace stan {
+namespace math {
+
+/** \ingroup prob_dists
+ * Returns the CDF of the geometric distribution. Given containers of
+ * matching sizes, returns the product of probabilities.
+ *
+ * The geometric distribution counts the number of failures before
+ * the first success: P(N <= n | theta) = 1 - (1 - theta)^(n + 1).
+ *
+ * @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::invalid_argument if container sizes mismatch
+ */
+template <typename T_n, typename T_prob,
+          require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
+              T_n, T_prob>* = nullptr>
+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_theta_ref = ref_type_t<T_prob>;
+  static constexpr const char* function = "geometric_cdf";
+  check_consistent_sizes(function, "Random variable", n,
+                         "Probability parameter", theta);
+  T_theta_ref theta_ref = theta;
+  const auto& n_arr = as_value_column_array_or_scalar(n);
+  const auto& theta_arr = as_value_column_array_or_scalar(theta_ref);
+  check_positive_finite(function, "Probability parameter", theta_arr);
+  check_less_or_equal(function, "Probability parameter", theta_arr, 1.0);
+
+  if (size_zero(n, theta)) {
+    return 1.0;
+  }
+
+  auto ops_partials = make_partials_propagator(theta_ref);
+
+  // P(N <= n) = 0 for n < 0
+  if (any(n_arr < 0)) {
+    return ops_partials.build(0.0);
+  }
+
+  // P_i = 1 - (1 - theta)^(n + 1) = -expm1((n + 1) * log1m(theta))
+  // For theta = 1: log1m(1) = -inf, (n+1)*-inf = -inf (n >= 0),
+  //   expm1(-inf) = -1, so P_i = 1 (correct: certain success means
+  //   N <= n always for n >= 0).
+  const auto& log1m_theta = log1m(theta_arr);
+  const auto& P_i = -expm1((n_arr + 1.0) * log1m_theta);
+  const T_partials_return P = prod(P_i);
+
+  if constexpr (is_autodiff_v<T_prob>) {
+    // d/dtheta P_i = (n + 1) * (1 - theta)^n
+    //             = (n + 1) * exp(n * log1m(theta))
+    // For n = 0: (n+1)*exp(0) = 1; the select avoids 0 * log1m(1) = NaN
+    //   when theta = 1.
+    // For n > 0, theta = 1: (n+1) * exp(n * -inf) = (n+1) * 0 = 0
+    //   (correct: derivative vanishes once CDF saturates at 1).
+    const auto& dP_dtheta = select(n_arr == 0, T_partials_return(1.0),
+                                   (n_arr + 1.0) * exp(n_arr * log1m_theta));
+    if constexpr (is_stan_scalar_v<T_prob>) {
+      partials<0>(ops_partials) = sum(P * elt_divide(dP_dtheta, P_i));
+    } else {
+      partials<0>(ops_partials) = P * elt_divide(dP_dtheta, P_i);
+    }
+  }
+
+  return ops_partials.build(P);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/prob/geometric_lccdf.hpp

@@ -0,0 +1,109 @@
+#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LCCDF_HPP
+#define STAN_MATH_PRIM_PROB_GEOMETRIC_LCCDF_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/any.hpp>
+#include <stan/math/prim/fun/as_value_column_array_or_scalar.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/inv.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/select.hpp>
+#include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/sum.hpp>
+#include <stan/math/prim/fun/value_of.hpp>
+#include <stan/math/prim/functor/partials_propagator.hpp>
+#include <limits>
+
+namespace stan {
+namespace math {
+
+/** \ingroup prob_dists
+ * Returns the log CCDF of the geometric distribution. Given containers of
+ * matching sizes, returns the log of the product of complementary
+ * probabilities.
+ *
+ * log P(N > n | theta) = log((1 - theta)^(n + 1)) = (n + 1) * log1m(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 product of complements
+ * @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,
+          require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
+              T_n, T_prob>* = nullptr>
+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_theta_ref = ref_type_t<T_prob>;
+  static constexpr const char* function = "geometric_lccdf";
+  check_consistent_sizes(function, "Random variable", n,
+                         "Probability parameter", theta);
+  T_n_ref n_ref = n;
+  T_theta_ref theta_ref = theta;
+  const auto& n_arr = as_value_column_array_or_scalar(n_ref);
+  const auto& theta_arr = as_value_column_array_or_scalar(theta_ref);
+  check_positive_finite(function, "Probability parameter", theta_arr);
+  check_less_or_equal(function, "Probability parameter", theta_arr, 1.0);
+
+  if (size_zero(n, theta)) {
+    return 0.0;
+  }
+
+  auto ops_partials = make_partials_propagator(theta_ref);
+
+  // n at INT_MAX: P(N > n) underflows to 0, lccdf = -inf.
+  // (The autodiff test framework probes the upper bound at INT_MAX,
+  // mirroring the early return used in neg_binomial_lccdf.)
+  if (any(n_arr == std::numeric_limits<int>::max())) {
+    return ops_partials.build(NEGATIVE_INFTY);
+  }
+
+  // theta = 1 with any n_i >= 0 makes (n_i + 1) * log1m(1) = -inf and
+  // the partials path divides by (theta - 1) = 0. Per-element loop here
+  // correctly pairs theta_i with n_i (e.g. theta=[0.5, 1.0], n=[2, -1]
+  // returns 3*log(0.5), not -inf). For theta_i = 1 paired with n_i < 0
+  // the per-element select below contributes 0 to both value and partials.
+  scalar_seq_view<T_n_ref> n_vec(n_ref);
+  scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
+  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[i] >= 0) {
+      return ops_partials.build(NEGATIVE_INFTY);
+    }
+  }
+
+  // Per-element: n_i < 0 contributes log(1) = 0 (since P(N > n_i) = 1
+  // for any n_i below the support), so they are no-ops in the sum.
+  // n_i >= 0 contributes (n_i + 1) * log1m(theta).
+  const auto& log1m_theta = log1m(theta_arr);
+  const auto& term
+      = select(n_arr < 0, T_partials_return(0), (n_arr + 1.0) * log1m_theta);
+  T_partials_return logP = sum(term);
+
+  if constexpr (is_autodiff_v<T_prob>) {
+    // d/dtheta of 0 = 0 for n_i < 0, and (n + 1) / (theta - 1) for n_i >= 0.
+    // theta = 1 was filtered above, so theta - 1 != 0 here.
+    const auto& dterm = select(n_arr < 0, T_partials_return(0),
+                               (n_arr + 1.0) * inv(theta_arr - 1.0));
+    if constexpr (is_stan_scalar_v<T_prob>) {
+      partials<0>(ops_partials) = sum(dterm);
+    } else {
+      partials<0>(ops_partials) = dterm;
+    }
+  }
+
+  return ops_partials.build(logP);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/prob/geometric_lcdf.hpp

@@ -0,0 +1,95 @@
+#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LCDF_HPP
+#define STAN_MATH_PRIM_PROB_GEOMETRIC_LCDF_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/any.hpp>
+#include <stan/math/prim/fun/as_value_column_array_or_scalar.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/elt_divide.hpp>
+#include <stan/math/prim/fun/exp.hpp>
+#include <stan/math/prim/fun/expm1.hpp>
+#include <stan/math/prim/fun/log1m.hpp>
+#include <stan/math/prim/fun/log1m_exp.hpp>
+#include <stan/math/prim/fun/select.hpp>
+#include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/sum.hpp>
+#include <stan/math/prim/fun/value_of.hpp>
+#include <stan/math/prim/functor/partials_propagator.hpp>
+
+namespace stan {
+namespace math {
+
+/** \ingroup prob_dists
+ * Returns the log CDF of the geometric distribution. Given containers of
+ * matching sizes, returns the log of the product of probabilities.
+ *
+ * log P(N <= n | theta) = log(1 - (1 - theta)^(n + 1))
+ *                       = log1m_exp((n + 1) * log1m(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::invalid_argument if container sizes mismatch
+ */
+template <typename T_n, typename T_prob,
+          require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
+              T_n, T_prob>* = nullptr>
+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_theta_ref = ref_type_t<T_prob>;
+  static constexpr const char* function = "geometric_lcdf";
+  check_consistent_sizes(function, "Random variable", n,
+                         "Probability parameter", theta);
+  T_theta_ref theta_ref = theta;
+  const auto& n_arr = as_value_column_array_or_scalar(n);
+  const auto& theta_arr = as_value_column_array_or_scalar(theta_ref);
+  check_positive_finite(function, "Probability parameter", theta_arr);
+  check_less_or_equal(function, "Probability parameter", theta_arr, 1.0);
+
+  if (size_zero(n, theta)) {
+    return 0.0;
+  }
+
+  auto ops_partials = make_partials_propagator(theta_ref);
+
+  // log P(N <= n) = -inf for n < 0
+  if (any(n_arr < 0)) {
+    return ops_partials.build(NEGATIVE_INFTY);
+  }
+
+  // log_q = log((1 - theta)^(n + 1)) = (n + 1) * log1m(theta)
+  // log P_i = log(1 - q_i) = log1m_exp(log_q_i)
+  // For theta = 1: log_q = -inf, log1m_exp(-inf) = log(1 - 0) = 0
+  //   (correct: P = 1 when success is certain).
+  const auto& log1m_theta = log1m(theta_arr);
+  const auto& log_q = (n_arr + 1.0) * log1m_theta;
+  const auto& log_P_i = log1m_exp(log_q);
+  T_partials_return logP = sum(log_P_i);
+
+  if constexpr (is_autodiff_v<T_prob>) {
+    // partial of log P_i = (dP_i/dtheta) / P_i
+    // dP_i/dtheta = (n + 1) * (1 - theta)^n = (n + 1) * exp(n * log1m_theta)
+    // P_i = -expm1(log_q)
+    // For n = 0: dP_dtheta = 1 (avoid 0 * log1m(1) = NaN at theta = 1).
+    // For n > 0, theta = 1: dP_dtheta = 0 and P_i = 1 -> partial = 0.
+    const auto& dP_dtheta = select(n_arr == 0, T_partials_return(1.0),
+                                   (n_arr + 1.0) * exp(n_arr * log1m_theta));
+    const auto& P_i = -expm1(log_q);
+    if constexpr (is_stan_scalar_v<T_prob>) {
+      partials<0>(ops_partials) = sum(elt_divide(dP_dtheta, P_i));
+    } else {
+      partials<0>(ops_partials) = elt_divide(dP_dtheta, P_i);
+    }
+  }
+
+  return ops_partials.build(logP);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif