From b61d325a81bc4d7ae556690a87a9c9ad71e5f595 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Tue, 24 Mar 2026 03:52:18 +0000 Subject: [PATCH 01/13] Add geometric distribution (lpmf, cdf, lcdf, lccdf, rng) Implements the geometric distribution as requested in #3098. The geometric distribution models the number of failures before the first success, with PMF: P(n|theta) = theta * (1-theta)^n where n in {0, 1, 2, ...} and theta in (0, 1]. This uses the number-of-failures parameterization, consistent with the geometric distribution being a special case of the negative binomial distribution with r=1 (neg_binomial(1, theta/(1-theta))). Verified: geometric_lpmf(n, theta) == neg_binomial_lpmf(n, 1, theta/(1-theta)) for all tested values. Includes: - geometric_lpmf: log probability mass function with autodiff - geometric_cdf: cumulative distribution function with autodiff - geometric_lcdf: log CDF with autodiff - geometric_lccdf: log complementary CDF with autodiff - geometric_rng: random number generation via inverse CDF method - Unit tests for all functions including distribution checks - CDF test fixture for gradient verification Co-Authored-By: Claude Opus 4.6 (1M context) --- stan/math/prim/prob.hpp | 5 + stan/math/prim/prob/geometric_cdf.hpp | 94 +++++++++++++++++ stan/math/prim/prob/geometric_lccdf.hpp | 83 +++++++++++++++ stan/math/prim/prob/geometric_lcdf.hpp | 88 ++++++++++++++++ stan/math/prim/prob/geometric_lpmf.hpp | 108 ++++++++++++++++++++ stan/math/prim/prob/geometric_rng.hpp | 76 ++++++++++++++ test/prob/geometric/geometric_cdf_test.hpp | 68 ++++++++++++ test/unit/math/prim/prob/geometric_test.cpp | 103 +++++++++++++++++++ 8 files changed, 625 insertions(+) create mode 100644 stan/math/prim/prob/geometric_cdf.hpp create mode 100644 stan/math/prim/prob/geometric_lccdf.hpp create mode 100644 stan/math/prim/prob/geometric_lcdf.hpp create mode 100644 stan/math/prim/prob/geometric_lpmf.hpp create mode 100644 stan/math/prim/prob/geometric_rng.hpp create mode 100644 test/prob/geometric/geometric_cdf_test.hpp create mode 100644 test/unit/math/prim/prob/geometric_test.cpp diff --git a/stan/math/prim/prob.hpp b/stan/math/prim/prob.hpp index 308b084a165..31a4dde94fa 100644 --- a/stan/math/prim/prob.hpp +++ b/stan/math/prim/prob.hpp @@ -111,6 +111,11 @@ #include #include #include +#include +#include +#include +#include +#include #include #include #include diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp new file mode 100644 index 00000000000..683b53a9185 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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. + * + * CDF: F(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 +inline return_type_t geometric_cdf(const T_n& n, + const T_prob& theta) { + using T_partials_return = partials_return_t; + using T_n_ref = ref_type_t; + using T_prob_ref = ref_type_t; + static constexpr const char* function = "geometric_cdf"; + + check_consistent_sizes(function, "Outcome variable", n, + "Success probability parameter", theta); + if (size_zero(n, theta)) { + return 1.0; + } + + T_n_ref n_ref = n; + T_prob_ref theta_ref = theta; + check_bounded(function, "Success probability parameter", + value_of(theta_ref), 0.0, 1.0); + + scalar_seq_view n_vec(n_ref); + scalar_seq_view 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) { + partials<0>(ops_partials)[i] + += np1 * ccdf_i / ((1.0 - theta_val) * cdf_i); + } + } + + if constexpr (is_autodiff_v) { + for (size_t i = 0; i < stan::math::size(theta); ++i) { + partials<0>(ops_partials)[i] *= cdf; + } + } + + return ops_partials.build(cdf); +} + +} // namespace math +} // namespace stan +#endif diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp new file mode 100644 index 00000000000..87e5318452d --- /dev/null +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -0,0 +1,83 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LCCDF_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_LCCDF_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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. + * + * log CCDF: (n + 1) * log(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::invalid_argument if container sizes mismatch + */ +template +inline return_type_t geometric_lccdf(const T_n& n, + const T_prob& theta) { + using T_partials_return = partials_return_t; + using T_n_ref = ref_type_t; + using T_prob_ref = ref_type_t; + static constexpr const char* function = "geometric_lccdf"; + + check_consistent_sizes(function, "Outcome variable", n, + "Success probability parameter", theta); + if (size_zero(n, theta)) { + return 0.0; + } + + T_n_ref n_ref = n; + T_prob_ref theta_ref = theta; + check_bounded(function, "Success probability parameter", + value_of(theta_ref), 0.0, 1.0); + + scalar_seq_view n_vec(n_ref); + scalar_seq_view 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) { + partials<0>(ops_partials)[i] += -np1 / (1.0 - theta_val); + } + } + + return ops_partials.build(log_ccdf); +} + +} // namespace math +} // namespace stan +#endif diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp new file mode 100644 index 00000000000..9653d8a28b1 --- /dev/null +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -0,0 +1,88 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LCDF_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_LCDF_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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. + * + * log CDF: log(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 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 +inline return_type_t geometric_lcdf(const T_n& n, + const T_prob& theta) { + using T_partials_return = partials_return_t; + using T_n_ref = ref_type_t; + using T_prob_ref = ref_type_t; + static constexpr const char* function = "geometric_lcdf"; + + check_consistent_sizes(function, "Outcome variable", n, + "Success probability parameter", theta); + if (size_zero(n, theta)) { + return 0.0; + } + + T_n_ref n_ref = n; + T_prob_ref theta_ref = theta; + check_bounded(function, "Success probability parameter", + value_of(theta_ref), 0.0, 1.0); + + scalar_seq_view n_vec(n_ref); + scalar_seq_view 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) { + const auto ccdf = stan::math::exp(log_ccdf); + partials<0>(ops_partials)[i] + += np1 * ccdf / ((1.0 - theta_val) * (1.0 - ccdf)); + } + } + + return ops_partials.build(log_cdf); +} + +} // namespace math +} // namespace stan +#endif diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp new file mode 100644 index 00000000000..945d24db0b0 --- /dev/null +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -0,0 +1,108 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LPMF_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_LPMF_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** \ingroup prob_dists + * Returns the log PMF of the geometric distribution. If containers + * of matching sizes are supplied, returns the log sum of probabilities. + * + * The geometric distribution with success probability \f$\theta\f$ has PMF + * \f[ + * P(n \mid \theta) = \theta (1 - \theta)^{n} + * \f] + * where \f$n \in \{0, 1, 2, \ldots\}\f$ and + * \f$\theta \in (0, 1]\f$. + * + * This is the number-of-failures parameterization, consistent with + * the geometric distribution being a special case of the negative + * binomial distribution with r = 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 log probability or log sum of probabilities + * @throw std::domain_error if theta is not in [0, 1] + * @throw std::domain_error if n is negative + * @throw std::invalid_argument if container sizes mismatch + */ +template * = nullptr> +inline return_type_t geometric_lpmf(const T_n& n, + const T_prob& theta) { + using std::log; + using T_partials_return = partials_return_t; + using T_n_ref = ref_type_t; + using T_prob_ref = ref_type_t; + static constexpr const char* function = "geometric_lpmf"; + check_consistent_sizes(function, "Outcome variable", n, + "Success probability parameter", theta); + if (size_zero(n, theta)) { + return 0.0; + } + + T_n_ref n_ref = n; + T_prob_ref theta_ref = theta; + check_nonnegative(function, "Outcome variable", n_ref); + check_bounded(function, "Success probability parameter", + value_of(theta_ref), 0.0, 1.0); + + if constexpr (!include_summand::value) { + return 0.0; + } + + auto ops_partials = make_partials_propagator(theta_ref); + + scalar_seq_view n_vec(n_ref); + scalar_seq_view theta_vec(theta_ref); + const size_t max_size_seq_view = max_size(n_ref, theta_ref); + T_partials_return logp(0.0); + + 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); + + // When theta == 1.0, P(n=0) = 1, P(n>0) = 0 + if (theta_val == 1.0) { + if (n_val > 0) { + return negative_infinity(); + } + // logp += 0 for n=0, theta=1 + continue; + } + + logp += log(theta_val) + n_val * log1m(theta_val); + + if constexpr (is_autodiff_v) { + partials<0>(ops_partials)[i] + += 1.0 / theta_val - n_val / (1.0 - theta_val); + } + } + return ops_partials.build(logp); +} + +template +inline return_type_t geometric_lpmf(const T_n& n, + const T_prob& theta) { + return geometric_lpmf(n, theta); +} + +} // namespace math +} // namespace stan +#endif diff --git a/stan/math/prim/prob/geometric_rng.hpp b/stan/math/prim/prob/geometric_rng.hpp new file mode 100644 index 00000000000..0e5a28be260 --- /dev/null +++ b/stan/math/prim/prob/geometric_rng.hpp @@ -0,0 +1,76 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_RNG_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_RNG_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** \ingroup prob_dists + * Return a geometric random variate with the given success probability + * parameter, using the given random number generator. + * + * theta can be a scalar or a one-dimensional container. + * + * The geometric distribution models the number of failures before the first + * success, with support on {0, 1, 2, ...}. This is the number-of-failures + * parameterization, consistent with the geometric distribution being a + * special case of the negative binomial with r = 1. + * + * Uses the inverse CDF method: n = floor(log(U) / log(1 - theta)) + * where U ~ Uniform(0, 1). + * + * @tparam T_prob type of success probability parameter + * @tparam RNG type of random number generator + * + * @param theta (Sequence of) success probability parameter(s) + * @param rng random number generator + * @return (Sequence of) geometric random variate(s) + * @throw std::domain_error if theta is not in (0, 1] + */ +template +inline auto geometric_rng(T_prob&& theta, RNG& rng) { + static constexpr const char* function = "geometric_rng"; + decltype(auto) theta_ref = to_ref(std::forward(theta)); + check_positive_finite(function, "Success probability parameter", + value_of(theta_ref)); + check_less_or_equal(function, "Success probability parameter", + value_of(theta_ref), 1.0); + + boost::random::uniform_real_distribution uniform(0.0, 1.0); + + if constexpr (is_stan_scalar_v) { + if (theta_ref == 1.0) { + return 0; + } + double u = uniform(rng); + return static_cast(std::floor(std::log(u) / log1m(theta_ref))); + } else { + auto theta_arr = as_array_or_scalar(theta_ref); + std::vector result(theta_arr.size()); + for (int i = 0; i < theta_arr.size(); i++) { + if (theta_arr[i] == 1.0) { + result[i] = 0; + } else { + double u_i = uniform(rng); + result[i] = static_cast( + std::floor(std::log(u_i) / log1m(theta_arr[i]))); + } + } + return result; + } +} + +} // namespace math +} // namespace stan +#endif diff --git a/test/prob/geometric/geometric_cdf_test.hpp b/test/prob/geometric/geometric_cdf_test.hpp new file mode 100644 index 00000000000..1a39cc85510 --- /dev/null +++ b/test/prob/geometric/geometric_cdf_test.hpp @@ -0,0 +1,68 @@ +// Arguments: Ints, Doubles +#include +#include + +using std::numeric_limits; +using std::vector; + +class AgradCdfGeometric : public AgradCdfTest { + public: + void valid_values(vector>& parameters, vector& cdf) { + vector param(2); + + // CDF(n=0|theta=0.5) = 1 - (0.5)^1 = 0.5 + param[0] = 0; // n + param[1] = 0.5; // theta + parameters.push_back(param); + cdf.push_back(0.5); + + // CDF(n=2|theta=0.5) = 1 - (0.5)^3 = 0.875 + param[0] = 2; // n + param[1] = 0.5; // theta + parameters.push_back(param); + cdf.push_back(0.875); + + // CDF(n=4|theta=0.3) = 1 - 0.7^5 = 0.83193 + param[0] = 4; // n + param[1] = 0.3; // theta + parameters.push_back(param); + cdf.push_back(0.83193); + + // CDF(n=0|theta=1.0) = 1.0 + param[0] = 0; // n + param[1] = 1.0; // theta + parameters.push_back(param); + cdf.push_back(1.0); + } + + void invalid_values(vector& index, vector& value) { + // theta + index.push_back(1U); + value.push_back(-0.1); + + index.push_back(1U); + value.push_back(1.1); + } + + bool has_lower_bound() { return false; } + + bool has_upper_bound() { return false; } + + template + stan::return_type_t cdf(const T_n& n, const T_prob& theta, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::geometric_cdf(n, theta); + } + + template + stan::return_type_t cdf_function(const T_n& n, + const T_prob& theta, + const T2&, const T3&, + const T4&, const T5&) { + using stan::math::log1m; + return 1.0 - stan::math::exp((n + 1.0) * log1m(theta)); + } +}; diff --git a/test/unit/math/prim/prob/geometric_test.cpp b/test/unit/math/prim/prob/geometric_test.cpp new file mode 100644 index 00000000000..4b0f537ca46 --- /dev/null +++ b/test/unit/math/prim/prob/geometric_test.cpp @@ -0,0 +1,103 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +class GeometricTestRig : public VectorIntRNGTestRig { + public: + GeometricTestRig() + : VectorIntRNGTestRig(10000, 10, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, + {0.1, 0.3, 0.5, 0.7, 0.9}, {1}, + {-0.1, 1.1}, {-1, 0, 2}) {} + + template + auto generate_samples(const T1& theta, const T2&, const T3&, + T_rng& rng) const { + return stan::math::geometric_rng(theta, rng); + } + + template + double pmf(int y, T1 theta, double, double) const { + if (y < 0) { + return 0.0; + } + return std::exp(stan::math::geometric_lpmf(y, theta)); + } +}; + +TEST(ProbDistributionsGeometric, errorCheck) { + check_dist_throws_all_types(GeometricTestRig()); +} + +TEST(ProbDistributionsGeometric, distributionCheck) { + check_counts_real(GeometricTestRig()); +} + +TEST(ProbDistributionsGeometric, error_check) { + boost::random::mt19937 rng; + + EXPECT_NO_THROW(stan::math::geometric_rng(0.5, rng)); + EXPECT_NO_THROW(stan::math::geometric_rng(0.1, rng)); + EXPECT_NO_THROW(stan::math::geometric_rng(0.9, rng)); + EXPECT_NO_THROW(stan::math::geometric_rng(1.0, rng)); + + EXPECT_THROW(stan::math::geometric_rng(0.0, rng), std::domain_error); + EXPECT_THROW(stan::math::geometric_rng(-0.5, rng), std::domain_error); + + EXPECT_THROW(stan::math::geometric_rng(stan::math::positive_infinity(), rng), + std::domain_error); + EXPECT_THROW(stan::math::geometric_rng(stan::math::not_a_number(), rng), + std::domain_error); +} + +TEST(ProbDistributionsGeometric, lpmf_values) { + // P(n=0|theta=0.5) = 0.5 + EXPECT_NEAR(stan::math::geometric_lpmf(0, 0.5), std::log(0.5), 1e-10); + // P(n=2|theta=0.5) = 0.5 * 0.5^2 = 0.125 + EXPECT_NEAR(stan::math::geometric_lpmf(2, 0.5), std::log(0.125), 1e-10); + // P(n=0|theta=1.0) = 1.0 + EXPECT_NEAR(stan::math::geometric_lpmf(0, 1.0), 0.0, 1e-10); + // P(n=1|theta=1.0) = 0 + EXPECT_EQ(stan::math::geometric_lpmf(1, 1.0), + stan::math::negative_infinity()); + // P(n=0|theta=0.3) = 0.3 + EXPECT_NEAR(stan::math::geometric_lpmf(0, 0.3), std::log(0.3), 1e-10); + // P(n=3|theta=0.3) = 0.3 * 0.7^3 = 0.1029 + EXPECT_NEAR(stan::math::geometric_lpmf(3, 0.3), std::log(0.1029), 1e-10); +} + +TEST(ProbDistributionsGeometric, cdf_values) { + // CDF(n=0|theta=0.5) = 1 - (1-0.5)^1 = 0.5 + EXPECT_NEAR(stan::math::geometric_cdf(0, 0.5), 0.5, 1e-10); + // CDF(n=2|theta=0.5) = 1 - (0.5)^3 = 0.875 + EXPECT_NEAR(stan::math::geometric_cdf(2, 0.5), 0.875, 1e-10); + // CDF(n=0|theta=1.0) = 1.0 + EXPECT_NEAR(stan::math::geometric_cdf(0, 1.0), 1.0, 1e-10); + // CDF(n=4|theta=0.3) = 1 - 0.7^5 = 0.83193 + EXPECT_NEAR(stan::math::geometric_cdf(4, 0.3), 0.83193, 1e-5); +} + +TEST(ProbDistributionsGeometric, cdf_below_support) { + EXPECT_DOUBLE_EQ(stan::math::geometric_cdf(-1, 0.5), 0.0); + EXPECT_DOUBLE_EQ(stan::math::geometric_lcdf(-1, 0.5), + stan::math::negative_infinity()); +} + +TEST(ProbDistributionsGeometric, lccdf_values) { + // CCDF(n=0|theta=0.5) = (1-0.5)^1 = 0.5 + EXPECT_NEAR(stan::math::geometric_lccdf(0, 0.5), std::log(0.5), 1e-10); + // CCDF(n=2|theta=0.5) = (0.5)^3 = 0.125 + EXPECT_NEAR(stan::math::geometric_lccdf(2, 0.5), std::log(0.125), 1e-10); +} + +TEST(ProbDistributionsGeometric, rng_theta_one) { + boost::random::mt19937 rng; + // theta=1.0 should always return 0 (always succeed on first trial) + for (int i = 0; i < 100; i++) { + EXPECT_EQ(stan::math::geometric_rng(1.0, rng), 0); + } +} From aa576fdd988848b4926fd1aaa2cc0089f5910fbb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Tue, 24 Mar 2026 04:25:03 +0000 Subject: [PATCH 02/13] Fix CodeRabbit review findings - Guard gradient computations against theta=0 division by zero in geometric_lpmf, geometric_cdf, and geometric_lcdf - Fix include ordering in prob.hpp (gaussian before geometric) - Fix signed/unsigned comparison in geometric_rng loop index Co-Authored-By: Claude Opus 4.6 (1M context) --- stan/math/prim/prob.hpp | 4 ++-- stan/math/prim/prob/geometric_cdf.hpp | 6 ++++-- stan/math/prim/prob/geometric_lcdf.hpp | 6 ++++-- stan/math/prim/prob/geometric_lpmf.hpp | 5 +++++ stan/math/prim/prob/geometric_rng.hpp | 2 +- 5 files changed, 16 insertions(+), 7 deletions(-) diff --git a/stan/math/prim/prob.hpp b/stan/math/prim/prob.hpp index 31a4dde94fa..33c39a183b8 100644 --- a/stan/math/prim/prob.hpp +++ b/stan/math/prim/prob.hpp @@ -111,13 +111,13 @@ #include #include #include +#include +#include #include #include #include #include #include -#include -#include #include #include #include diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index 683b53a9185..bcd33fac0ff 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -75,8 +75,10 @@ inline return_type_t geometric_cdf(const T_n& n, cdf *= cdf_i; if constexpr (is_autodiff_v) { - partials<0>(ops_partials)[i] - += np1 * ccdf_i / ((1.0 - theta_val) * cdf_i); + if (cdf_i > 0.0 && theta_val > 0.0) { + partials<0>(ops_partials)[i] + += np1 * ccdf_i / ((1.0 - theta_val) * cdf_i); + } } } diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index 9653d8a28b1..8be35548885 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -75,8 +75,10 @@ inline return_type_t geometric_lcdf(const T_n& n, if constexpr (is_autodiff_v) { const auto ccdf = stan::math::exp(log_ccdf); - partials<0>(ops_partials)[i] - += np1 * ccdf / ((1.0 - theta_val) * (1.0 - ccdf)); + if (theta_val > 0.0 && ccdf < 1.0) { + partials<0>(ops_partials)[i] + += np1 * ccdf / ((1.0 - theta_val) * (1.0 - ccdf)); + } } } diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index 945d24db0b0..a802d0be353 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -78,6 +78,11 @@ inline return_type_t geometric_lpmf(const T_n& n, const auto theta_val = theta_vec.val(i); const auto n_val = n_vec.val(i); + // When theta == 0.0, P(n) = 0 for all n + if (theta_val == 0.0) { + return negative_infinity(); + } + // When theta == 1.0, P(n=0) = 1, P(n>0) = 0 if (theta_val == 1.0) { if (n_val > 0) { diff --git a/stan/math/prim/prob/geometric_rng.hpp b/stan/math/prim/prob/geometric_rng.hpp index 0e5a28be260..efd1465f292 100644 --- a/stan/math/prim/prob/geometric_rng.hpp +++ b/stan/math/prim/prob/geometric_rng.hpp @@ -58,7 +58,7 @@ inline auto geometric_rng(T_prob&& theta, RNG& rng) { } else { auto theta_arr = as_array_or_scalar(theta_ref); std::vector result(theta_arr.size()); - for (int i = 0; i < theta_arr.size(); i++) { + for (size_t i = 0; i < theta_arr.size(); i++) { if (theta_arr[i] == 1.0) { result[i] = 0; } else { From 8f302419a898e1bfc35e6e2da9845cd3c12e12db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Wed, 25 Mar 2026 03:34:51 +0000 Subject: [PATCH 03/13] 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. --- stan/math/prim/prob/geometric_cdf.hpp | 66 +++++++-------- stan/math/prim/prob/geometric_lccdf.hpp | 53 ++++++------ stan/math/prim/prob/geometric_lcdf.hpp | 64 +++++++-------- stan/math/prim/prob/geometric_lpmf.hpp | 82 ++++++++----------- stan/math/prim/prob/geometric_rng.hpp | 27 ++---- .../geometric/geometric_ccdf_log_test.hpp | 63 ++++++++++++++ .../prob/geometric/geometric_cdf_log_test.hpp | 64 +++++++++++++++ test/prob/geometric/geometric_cdf_test.hpp | 6 -- test/prob/geometric/geometric_test.hpp | 73 +++++++++++++++++ 9 files changed, 333 insertions(+), 165 deletions(-) create mode 100644 test/prob/geometric/geometric_ccdf_log_test.hpp create mode 100644 test/prob/geometric/geometric_cdf_log_test.hpp create mode 100644 test/prob/geometric/geometric_test.hpp diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index bcd33fac0ff..bb25af3168f 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -3,26 +3,24 @@ #include #include -#include -#include -#include #include #include #include #include #include -#include -#include +#include +#include +#include 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 @@ -30,13 +28,12 @@ namespace math { * @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 inline return_type_t geometric_cdf(const T_n& n, const T_prob& theta) { - using T_partials_return = partials_return_t; using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_cdf"; @@ -53,42 +50,39 @@ inline return_type_t geometric_cdf(const T_n& n, value_of(theta_ref), 0.0, 1.0); scalar_seq_view n_vec(n_ref); - scalar_seq_view 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) { - 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 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) { - for (size_t i = 0; i < stan::math::size(theta); ++i) { - partials<0>(ops_partials)[i] *= cdf; + if constexpr (is_stan_scalar_v) { + const auto beta = theta_ref / (1.0 - theta_ref); + return neg_binomial_cdf(n_ref, 1, beta); + } else if constexpr (is_std_vector_v) { + std::vector> 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 diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index 87e5318452d..a9c41ed08c4 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -3,24 +3,24 @@ #include #include -#include -#include #include #include #include #include #include -#include +#include +#include +#include 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 @@ -28,13 +28,12 @@ namespace math { * @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 inline return_type_t geometric_lccdf(const T_n& n, const T_prob& theta) { - using T_partials_return = partials_return_t; using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_lccdf"; @@ -51,31 +50,35 @@ inline return_type_t geometric_lccdf(const T_n& n, value_of(theta_ref), 0.0, 1.0); scalar_seq_view n_vec(n_ref); - scalar_seq_view 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) { - partials<0>(ops_partials)[i] += -np1 / (1.0 - theta_val); + // theta = 1 => CCDF = 0 for n >= 0, log CCDF = -inf + scalar_seq_view 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) { + const auto beta = theta_ref / (1.0 - theta_ref); + return neg_binomial_lccdf(n_ref, 1, beta); + } else if constexpr (is_std_vector_v) { + std::vector> 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 diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index 8be35548885..bf2e17a7bc3 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -3,26 +3,24 @@ #include #include -#include -#include -#include -#include #include #include #include #include #include -#include +#include +#include +#include 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 @@ -30,13 +28,12 @@ namespace math { * @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 inline return_type_t geometric_lcdf(const T_n& n, const T_prob& theta) { - using T_partials_return = partials_return_t; using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_lcdf"; @@ -53,36 +50,39 @@ inline return_type_t geometric_lcdf(const T_n& n, value_of(theta_ref), 0.0, 1.0); scalar_seq_view n_vec(n_ref); - scalar_seq_view 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) { - 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 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) { + const auto beta = theta_ref / (1.0 - theta_ref); + return neg_binomial_lcdf(n_ref, 1, beta); + } else if constexpr (is_std_vector_v) { + std::vector> 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 diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index a802d0be353..2718d78062e 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -3,15 +3,14 @@ #include #include -#include -#include -#include #include #include #include #include #include -#include +#include +#include +#include namespace stan { namespace math { @@ -20,16 +19,8 @@ namespace math { * Returns the log PMF of the geometric distribution. If containers * of matching sizes are supplied, returns the log sum of probabilities. * - * The geometric distribution with success probability \f$\theta\f$ has PMF - * \f[ - * P(n \mid \theta) = \theta (1 - \theta)^{n} - * \f] - * where \f$n \in \{0, 1, 2, \ldots\}\f$ and - * \f$\theta \in (0, 1]\f$. - * - * This is the number-of-failures parameterization, consistent with - * the geometric distribution being a special case of the negative - * binomial distribution with r = 1. + * The geometric distribution is a special case of the negative + * binomial distribution with alpha = 1 and beta = theta / (1 - theta). * * @tparam T_n type of outcome variable * @tparam T_prob type of success probability parameter @@ -37,7 +28,7 @@ namespace math { * @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::domain_error if n is negative * @throw std::invalid_argument if container sizes mismatch */ @@ -46,8 +37,6 @@ template * = nullptr> inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { - using std::log; - using T_partials_return = partials_return_t; using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_lpmf"; @@ -63,43 +52,44 @@ inline return_type_t geometric_lpmf(const T_n& n, check_bounded(function, "Success probability parameter", value_of(theta_ref), 0.0, 1.0); - if constexpr (!include_summand::value) { - return 0.0; - } - - auto ops_partials = make_partials_propagator(theta_ref); - + // theta = 1 => deterministic: P(0) = 1, P(n>0) = 0 + // Cannot delegate since beta = theta / (1 - theta) diverges scalar_seq_view n_vec(n_ref); scalar_seq_view theta_vec(theta_ref); - const size_t max_size_seq_view = max_size(n_ref, theta_ref); - T_partials_return logp(0.0); - - 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); - - // When theta == 0.0, P(n) = 0 for all n - if (theta_val == 0.0) { - return negative_infinity(); - } - - // When theta == 1.0, P(n=0) = 1, P(n>0) = 0 - if (theta_val == 1.0) { - if (n_val > 0) { + 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) { + if (n_vec[i] > 0) { return negative_infinity(); } - // logp += 0 for n=0, theta=1 - continue; } + } + 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; + } - logp += log(theta_val) + n_val * log1m(theta_val); - - if constexpr (is_autodiff_v) { - partials<0>(ops_partials)[i] - += 1.0 / theta_val - n_val / (1.0 - theta_val); + // geometric(theta) = neg_binomial(1, theta / (1 - theta)) + if constexpr (is_stan_scalar_v) { + const auto beta = theta_ref / (1.0 - theta_ref); + return neg_binomial_lpmf(n_ref, 1, beta); + } else if constexpr (is_std_vector_v) { + std::vector> 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_lpmf(n_ref, 1, beta); + } else { + const auto beta = elt_divide(theta_ref, subtract(1.0, theta_ref)); + return neg_binomial_lpmf(n_ref, 1, beta); } - return ops_partials.build(logp); } template diff --git a/stan/math/prim/prob/geometric_rng.hpp b/stan/math/prim/prob/geometric_rng.hpp index efd1465f292..fa7619a86c6 100644 --- a/stan/math/prim/prob/geometric_rng.hpp +++ b/stan/math/prim/prob/geometric_rng.hpp @@ -4,12 +4,9 @@ #include #include #include -#include -#include #include #include -#include -#include +#include #include #include @@ -20,15 +17,8 @@ namespace math { * Return a geometric random variate with the given success probability * parameter, using the given random number generator. * - * theta can be a scalar or a one-dimensional container. - * - * The geometric distribution models the number of failures before the first - * success, with support on {0, 1, 2, ...}. This is the number-of-failures - * parameterization, consistent with the geometric distribution being a - * special case of the negative binomial with r = 1. - * - * Uses the inverse CDF method: n = floor(log(U) / log(1 - theta)) - * where U ~ Uniform(0, 1). + * Delegates to neg_binomial_rng with alpha = 1 and + * beta = theta / (1 - theta), with a special case for theta = 1. * * @tparam T_prob type of success probability parameter * @tparam RNG type of random number generator @@ -47,14 +37,12 @@ inline auto geometric_rng(T_prob&& theta, RNG& rng) { check_less_or_equal(function, "Success probability parameter", value_of(theta_ref), 1.0); - boost::random::uniform_real_distribution uniform(0.0, 1.0); - if constexpr (is_stan_scalar_v) { if (theta_ref == 1.0) { return 0; } - double u = uniform(rng); - return static_cast(std::floor(std::log(u) / log1m(theta_ref))); + double beta = theta_ref / (1.0 - theta_ref); + return neg_binomial_rng(1, beta, rng); } else { auto theta_arr = as_array_or_scalar(theta_ref); std::vector result(theta_arr.size()); @@ -62,9 +50,8 @@ inline auto geometric_rng(T_prob&& theta, RNG& rng) { if (theta_arr[i] == 1.0) { result[i] = 0; } else { - double u_i = uniform(rng); - result[i] = static_cast( - std::floor(std::log(u_i) / log1m(theta_arr[i]))); + double beta_i = theta_arr[i] / (1.0 - theta_arr[i]); + result[i] = neg_binomial_rng(1, beta_i, rng); } } return result; diff --git a/test/prob/geometric/geometric_ccdf_log_test.hpp b/test/prob/geometric/geometric_ccdf_log_test.hpp new file mode 100644 index 00000000000..6a5595c47d7 --- /dev/null +++ b/test/prob/geometric/geometric_ccdf_log_test.hpp @@ -0,0 +1,63 @@ +// Arguments: Ints, Doubles +#include +#include + +using std::numeric_limits; +using std::vector; + +class AgradCcdfLogGeometric : public AgradCcdfLogTest { + public: + void valid_values(vector>& parameters, + vector& ccdf_log) { + vector param(2); + + // lccdf(n=0|theta=0.5) = log(0.5) = -0.693147... + param[0] = 0; // n + param[1] = 0.5; // theta + parameters.push_back(param); + ccdf_log.push_back(-0.69314718055994528623); + + // lccdf(n=2|theta=0.5) = log(0.125) = -2.079441... + param[0] = 2; // n + param[1] = 0.5; // theta + parameters.push_back(param); + ccdf_log.push_back(-2.07944154167983574766); + + // lccdf(n=4|theta=0.3) = 5*log(0.7) = -1.783374... + param[0] = 4; // n + param[1] = 0.3; // theta + parameters.push_back(param); + ccdf_log.push_back(-1.78337471969366179181); + } + + void invalid_values(vector& index, vector& value) { + // theta + index.push_back(1U); + value.push_back(-0.1); + + index.push_back(1U); + value.push_back(1.1); + } + + bool has_lower_bound() { return false; } + + bool has_upper_bound() { return false; } + + template + stan::return_type_t ccdf_log(const T_n& n, const T_prob& theta, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::geometric_lccdf(n, theta); + } + + template + stan::return_type_t ccdf_log_function(const T_n& n, + const T_prob& theta, + const T2&, const T3&, + const T4&, const T5&) { + using stan::math::log1m; + return (n + 1.0) * log1m(theta); + } +}; diff --git a/test/prob/geometric/geometric_cdf_log_test.hpp b/test/prob/geometric/geometric_cdf_log_test.hpp new file mode 100644 index 00000000000..9843a30c3c0 --- /dev/null +++ b/test/prob/geometric/geometric_cdf_log_test.hpp @@ -0,0 +1,64 @@ +// Arguments: Ints, Doubles +#include +#include + +using std::numeric_limits; +using std::vector; + +class AgradCdfLogGeometric : public AgradCdfLogTest { + public: + void valid_values(vector>& parameters, + vector& cdf_log) { + vector param(2); + + // lcdf(n=0|theta=0.5) = log(0.5) = -0.693147... + param[0] = 0; // n + param[1] = 0.5; // theta + parameters.push_back(param); + cdf_log.push_back(-0.69314718055994528623); + + // lcdf(n=2|theta=0.5) = log(0.875) = -0.133531... + param[0] = 2; // n + param[1] = 0.5; // theta + parameters.push_back(param); + cdf_log.push_back(-0.13353139262452262681); + + // lcdf(n=4|theta=0.3) = log(0.83193) = -0.184006... + param[0] = 4; // n + param[1] = 0.3; // theta + parameters.push_back(param); + cdf_log.push_back(-0.18400697631582843550); + } + + void invalid_values(vector& index, vector& value) { + // theta + index.push_back(1U); + value.push_back(-0.1); + + index.push_back(1U); + value.push_back(1.1); + } + + bool has_lower_bound() { return false; } + + bool has_upper_bound() { return false; } + + template + stan::return_type_t cdf_log(const T_n& n, const T_prob& theta, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::geometric_lcdf(n, theta); + } + + template + stan::return_type_t cdf_log_function(const T_n& n, + const T_prob& theta, + const T2&, const T3&, + const T4&, const T5&) { + using stan::math::log1m; + using stan::math::log1m_exp; + return log1m_exp((n + 1.0) * log1m(theta)); + } +}; diff --git a/test/prob/geometric/geometric_cdf_test.hpp b/test/prob/geometric/geometric_cdf_test.hpp index 1a39cc85510..a2966a9a65f 100644 --- a/test/prob/geometric/geometric_cdf_test.hpp +++ b/test/prob/geometric/geometric_cdf_test.hpp @@ -27,12 +27,6 @@ class AgradCdfGeometric : public AgradCdfTest { param[1] = 0.3; // theta parameters.push_back(param); cdf.push_back(0.83193); - - // CDF(n=0|theta=1.0) = 1.0 - param[0] = 0; // n - param[1] = 1.0; // theta - parameters.push_back(param); - cdf.push_back(1.0); } void invalid_values(vector& index, vector& value) { diff --git a/test/prob/geometric/geometric_test.hpp b/test/prob/geometric/geometric_test.hpp new file mode 100644 index 00000000000..c697637d714 --- /dev/null +++ b/test/prob/geometric/geometric_test.hpp @@ -0,0 +1,73 @@ +// Arguments: Ints, Doubles +#include +#include + +using stan::math::var; +using std::numeric_limits; +using std::vector; + +class AgradDistributionsGeometric : public AgradDistributionTest { + public: + void valid_values(vector>& parameters, + vector& log_prob) { + vector param(2); + + // lpmf(n=0|theta=0.5) = log(0.5) = -0.693147... + param[0] = 0; // n + param[1] = 0.5; // theta + parameters.push_back(param); + log_prob.push_back(-0.69314718055994528623); + + // lpmf(n=2|theta=0.5) = log(0.5) + 2*log(0.5) = -2.079441... + param[0] = 2; // n + param[1] = 0.5; // theta + parameters.push_back(param); + log_prob.push_back(-2.07944154167983574766); + + // lpmf(n=4|theta=0.3) = log(0.3) + 4*log(0.7) = -2.630672... + param[0] = 4; // n + param[1] = 0.3; // theta + parameters.push_back(param); + log_prob.push_back(-2.63067258008086568566); + } + + void invalid_values(vector& index, vector& value) { + // n + index.push_back(0U); + value.push_back(-1); + + // theta + index.push_back(1U); + value.push_back(-0.1); + + index.push_back(1U); + value.push_back(1.1); + } + + template + stan::return_type_t log_prob(const T_n& n, const T_prob& theta, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::geometric_lpmf(n, theta); + } + + template + stan::return_type_t log_prob(const T_n& n, const T_prob& theta, + const T2&, const T3&, const T4&, + const T5&) { + return stan::math::geometric_lpmf(n, theta); + } + + template + stan::return_type_t log_prob_function(const T_n& n, + const T_prob& theta, + const T2&, const T3&, + const T4&, const T5&) { + using stan::math::log1m; + using std::log; + return log(theta) + n * log1m(theta); + } +}; From 8485f4e083ab806e239e8ea14476fd8bfcf88199 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Wed, 25 Mar 2026 03:37:21 +0000 Subject: [PATCH 04/13] Add missing #include to geometric distribution files --- stan/math/prim/prob/geometric_cdf.hpp | 1 + stan/math/prim/prob/geometric_lccdf.hpp | 1 + stan/math/prim/prob/geometric_lcdf.hpp | 1 + stan/math/prim/prob/geometric_lpmf.hpp | 1 + 4 files changed, 4 insertions(+) diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index bb25af3168f..34f29e7b32a 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index a9c41ed08c4..ba9b3398395 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index bf2e17a7bc3..41053c4f03c 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index 2718d78062e..7337064fe3e 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include From 88d6b53d09db1ffe86daad3ffc41b4adb0386c6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Wed, 25 Mar 2026 16:25:32 +0000 Subject: [PATCH 05/13] Merge theta=1.0 check into single pass in geometric_lpmf --- stan/math/prim/prob/geometric_lpmf.hpp | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index 7337064fe3e..015a83e541a 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -58,18 +58,14 @@ inline return_type_t geometric_lpmf(const T_n& n, scalar_seq_view n_vec(n_ref); scalar_seq_view theta_vec(theta_ref); const size_t max_sz = max_size(n_ref, theta_ref); + bool all_theta_one = true; for (size_t i = 0; i < max_sz; i++) { if (value_of(theta_vec[i]) == 1.0) { if (n_vec[i] > 0) { return negative_infinity(); } - } - } - bool all_theta_one = true; - for (size_t i = 0; i < stan::math::size(theta); i++) { - if (value_of(theta_vec[i]) != 1.0) { + } else { all_theta_one = false; - break; } } if (all_theta_one) { From 84319aea1535c9e0678991105db179bdcb4481e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Wed, 25 Mar 2026 16:32:16 +0000 Subject: [PATCH 06/13] Apply clang-format to geometric distribution files --- stan/math/prim/prob/geometric_cdf.hpp | 7 +++---- stan/math/prim/prob/geometric_lccdf.hpp | 4 ++-- stan/math/prim/prob/geometric_lcdf.hpp | 7 +++---- stan/math/prim/prob/geometric_lpmf.hpp | 10 ++++------ 4 files changed, 12 insertions(+), 16 deletions(-) diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index 34f29e7b32a..5e0c2474f67 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -33,8 +33,7 @@ namespace math { * @throw std::invalid_argument if container sizes mismatch */ template -inline return_type_t geometric_cdf(const T_n& n, - const T_prob& theta) { +inline return_type_t geometric_cdf(const T_n& n, const T_prob& theta) { using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_cdf"; @@ -47,8 +46,8 @@ inline return_type_t geometric_cdf(const T_n& n, T_n_ref n_ref = n; T_prob_ref theta_ref = theta; - check_bounded(function, "Success probability parameter", - value_of(theta_ref), 0.0, 1.0); + check_bounded(function, "Success probability parameter", value_of(theta_ref), + 0.0, 1.0); scalar_seq_view n_vec(n_ref); for (int i = 0; i < stan::math::size(n); i++) { diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index ba9b3398395..9927caaf801 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -47,8 +47,8 @@ inline return_type_t geometric_lccdf(const T_n& n, T_n_ref n_ref = n; T_prob_ref theta_ref = theta; - check_bounded(function, "Success probability parameter", - value_of(theta_ref), 0.0, 1.0); + check_bounded(function, "Success probability parameter", value_of(theta_ref), + 0.0, 1.0); scalar_seq_view n_vec(n_ref); for (int i = 0; i < stan::math::size(n); i++) { diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index 41053c4f03c..0032d0c7ec4 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -33,8 +33,7 @@ namespace math { * @throw std::invalid_argument if container sizes mismatch */ template -inline return_type_t geometric_lcdf(const T_n& n, - const T_prob& theta) { +inline return_type_t geometric_lcdf(const T_n& n, const T_prob& theta) { using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_lcdf"; @@ -47,8 +46,8 @@ inline return_type_t geometric_lcdf(const T_n& n, T_n_ref n_ref = n; T_prob_ref theta_ref = theta; - check_bounded(function, "Success probability parameter", - value_of(theta_ref), 0.0, 1.0); + check_bounded(function, "Success probability parameter", value_of(theta_ref), + 0.0, 1.0); scalar_seq_view n_vec(n_ref); for (int i = 0; i < stan::math::size(n); i++) { diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index 015a83e541a..df0a8f2a03e 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -36,8 +36,7 @@ namespace math { template * = nullptr> -inline return_type_t geometric_lpmf(const T_n& n, - const T_prob& theta) { +inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { using T_n_ref = ref_type_t; using T_prob_ref = ref_type_t; static constexpr const char* function = "geometric_lpmf"; @@ -50,8 +49,8 @@ inline return_type_t geometric_lpmf(const T_n& n, T_n_ref n_ref = n; T_prob_ref theta_ref = theta; check_nonnegative(function, "Outcome variable", n_ref); - check_bounded(function, "Success probability parameter", - value_of(theta_ref), 0.0, 1.0); + check_bounded(function, "Success probability parameter", value_of(theta_ref), + 0.0, 1.0); // theta = 1 => deterministic: P(0) = 1, P(n>0) = 0 // Cannot delegate since beta = theta / (1 - theta) diverges @@ -90,8 +89,7 @@ inline return_type_t geometric_lpmf(const T_n& n, } template -inline return_type_t geometric_lpmf(const T_n& n, - const T_prob& theta) { +inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { return geometric_lpmf(n, theta); } From 707eba37e4cd23b74a158c8b15a5d60913e03683 Mon Sep 17 00:00:00 2001 From: Stan Jenkins Date: Wed, 25 Mar 2026 12:33:39 -0400 Subject: [PATCH 07/13] [Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1 --- test/prob/geometric/geometric_cdf_test.hpp | 6 +++--- test/prob/geometric/geometric_test.hpp | 14 +++++++------- test/unit/math/prim/prob/geometric_test.cpp | 4 ++-- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/test/prob/geometric/geometric_cdf_test.hpp b/test/prob/geometric/geometric_cdf_test.hpp index a2966a9a65f..cfd36f66c4d 100644 --- a/test/prob/geometric/geometric_cdf_test.hpp +++ b/test/prob/geometric/geometric_cdf_test.hpp @@ -53,9 +53,9 @@ class AgradCdfGeometric : public AgradCdfTest { template stan::return_type_t cdf_function(const T_n& n, - const T_prob& theta, - const T2&, const T3&, - const T4&, const T5&) { + const T_prob& theta, const T2&, + const T3&, const T4&, + const T5&) { using stan::math::log1m; return 1.0 - stan::math::exp((n + 1.0) * log1m(theta)); } diff --git a/test/prob/geometric/geometric_test.hpp b/test/prob/geometric/geometric_test.hpp index c697637d714..8acf08fd22f 100644 --- a/test/prob/geometric/geometric_test.hpp +++ b/test/prob/geometric/geometric_test.hpp @@ -47,25 +47,25 @@ class AgradDistributionsGeometric : public AgradDistributionTest { template stan::return_type_t log_prob(const T_n& n, const T_prob& theta, - const T2&, const T3&, const T4&, - const T5&) { + const T2&, const T3&, const T4&, + const T5&) { return stan::math::geometric_lpmf(n, theta); } template stan::return_type_t log_prob(const T_n& n, const T_prob& theta, - const T2&, const T3&, const T4&, - const T5&) { + const T2&, const T3&, const T4&, + const T5&) { return stan::math::geometric_lpmf(n, theta); } template stan::return_type_t log_prob_function(const T_n& n, - const T_prob& theta, - const T2&, const T3&, - const T4&, const T5&) { + const T_prob& theta, const T2&, + const T3&, const T4&, + const T5&) { using stan::math::log1m; using std::log; return log(theta) + n * log1m(theta); diff --git a/test/unit/math/prim/prob/geometric_test.cpp b/test/unit/math/prim/prob/geometric_test.cpp index 4b0f537ca46..e564f78cfbf 100644 --- a/test/unit/math/prim/prob/geometric_test.cpp +++ b/test/unit/math/prim/prob/geometric_test.cpp @@ -11,8 +11,8 @@ class GeometricTestRig : public VectorIntRNGTestRig { public: GeometricTestRig() : VectorIntRNGTestRig(10000, 10, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, - {0.1, 0.3, 0.5, 0.7, 0.9}, {1}, - {-0.1, 1.1}, {-1, 0, 2}) {} + {0.1, 0.3, 0.5, 0.7, 0.9}, {1}, {-0.1, 1.1}, + {-1, 0, 2}) {} template auto generate_samples(const T1& theta, const T2&, const T3&, From 640cb1a1cf6659c1a6b8b875bd9cb3d7a146a332 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Sat, 25 Apr 2026 01:28:32 +0000 Subject: [PATCH 08/13] Refactor geometric distribution to use partials_propagator Replace neg_binomial_* delegation with a full autodiff implementation per PR #3299 review. All four functions compute analytic partials on Eigen arrays directly, with explicit boundary handling for theta=0, theta=1, n=0, n=INT_MAX. Tests: 180/180 autodiff + 8/8 prim pass. --- stan/math/prim/prob/geometric_cdf.hpp | 103 ++++++++++++----------- stan/math/prim/prob/geometric_lccdf.hpp | 101 ++++++++++++---------- stan/math/prim/prob/geometric_lcdf.hpp | 106 +++++++++++++----------- stan/math/prim/prob/geometric_lpmf.hpp | 93 +++++++++++---------- 4 files changed, 221 insertions(+), 182 deletions(-) diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index 5e0c2474f67..17f38648d0f 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -3,15 +3,19 @@ #include #include -#include -#include -#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include #include -#include -#include -#include -#include +#include namespace stan { namespace math { @@ -20,8 +24,8 @@ namespace math { * Returns the CDF of the geometric distribution. Given containers of * matching sizes, returns the product of probabilities. * - * Delegates to the negative binomial CDF with alpha = 1 and - * beta = theta / (1 - theta). + * 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 @@ -29,60 +33,65 @@ namespace math { * @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 +template * = nullptr> inline return_type_t geometric_cdf(const T_n& n, const T_prob& theta) { - using T_n_ref = ref_type_t; - using T_prob_ref = ref_type_t; + using T_partials_return = partials_return_t; + using T_theta_ref = ref_type_t; 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_bounded(function, "Probability parameter", theta_arr, 0.0, 1.0); - check_consistent_sizes(function, "Outcome variable", n, - "Success probability parameter", theta); if (size_zero(n, theta)) { return 1.0; } - T_n_ref n_ref = n; - T_prob_ref theta_ref = theta; - check_bounded(function, "Success probability parameter", value_of(theta_ref), - 0.0, 1.0); + auto ops_partials = make_partials_propagator(theta_ref); - scalar_seq_view n_vec(n_ref); - for (int i = 0; i < stan::math::size(n); i++) { - if (n_vec.val(i) < 0) { - return 0.0; - } + // P(N <= n) = 0 for n < 0 + if (any(n_arr < 0)) { + return ops_partials.build(0.0); } - // theta = 1 => CDF is always 1 for n >= 0 - scalar_seq_view 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; + // theta = 0 is degenerate: P(N <= n) = 0 for any finite n. + // Avoid divide-by-zero in the partials path below. + if (any(theta_arr == 0.0)) { + return ops_partials.build(0.0); } - if constexpr (is_stan_scalar_v) { - const auto beta = theta_ref / (1.0 - theta_ref); - return neg_binomial_cdf(n_ref, 1, beta); - } else if constexpr (is_std_vector_v) { - std::vector> 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])); + // 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) { + // 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) { + 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 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(P); } } // namespace math diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index 9927caaf801..2532a59c2d4 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -3,83 +3,94 @@ #include #include -#include -#include -#include +#include +#include +#include +#include +#include #include +#include #include -#include -#include -#include -#include +#include +#include namespace stan { namespace math { /** \ingroup prob_dists * Returns the log CCDF of the geometric distribution. Given containers of - * matching sizes, returns the log sum of probabilities. + * matching sizes, returns the log of the product of complementary + * probabilities. * - * Delegates to the negative binomial log CCDF with alpha = 1 and - * beta = theta / (1 - theta). + * 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 sum - * @throw std::domain_error if theta is not in (0, 1] + * @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 +template * = nullptr> inline return_type_t geometric_lccdf(const T_n& n, const T_prob& theta) { - using T_n_ref = ref_type_t; - using T_prob_ref = ref_type_t; + using T_partials_return = partials_return_t; + using T_theta_ref = ref_type_t; static constexpr const char* function = "geometric_lccdf"; + 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_bounded(function, "Probability parameter", theta_arr, 0.0, 1.0); - check_consistent_sizes(function, "Outcome variable", n, - "Success probability parameter", theta); if (size_zero(n, theta)) { return 0.0; } - T_n_ref n_ref = n; - T_prob_ref theta_ref = theta; - check_bounded(function, "Success probability parameter", value_of(theta_ref), - 0.0, 1.0); + auto ops_partials = make_partials_propagator(theta_ref); - scalar_seq_view n_vec(n_ref); - for (int i = 0; i < stan::math::size(n); i++) { - if (n_vec.val(i) < 0) { - return 0.0; - } + // log P(N > n) = 0 (i.e. P = 1) when n < 0, matching the existing + // implementation that short-circuits on the first negative element. + if (any(n_arr < 0)) { + return ops_partials.build(0.0); } - // theta = 1 => CCDF = 0 for n >= 0, log CCDF = -inf - scalar_seq_view 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(); - } + // 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::max())) { + return ops_partials.build(NEGATIVE_INFTY); + } + + // theta = 1 means certain success, so P(N > n) = 0 for n >= 0 and the + // log is -inf. The partials path divides by (theta - 1) = 0, so we + // short-circuit. + if (any(theta_arr == 1.0)) { + return ops_partials.build(NEGATIVE_INFTY); } - if constexpr (is_stan_scalar_v) { - const auto beta = theta_ref / (1.0 - theta_ref); - return neg_binomial_lccdf(n_ref, 1, beta); - } else if constexpr (is_std_vector_v) { - std::vector> 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])); + // log P(N > n) = (n + 1) * log1m(theta) + // For theta = 0: log1m(0) = 0, lccdf = 0 (correct: certain failure). + const auto& log1m_theta = log1m(theta_arr); + T_partials_return logP = sum((n_arr + 1.0) * log1m_theta); + + if constexpr (is_autodiff_v) { + // d/dtheta (n + 1) * log1m(theta) = -(n + 1) / (1 - theta) + // = (n + 1) / (theta - 1) + // theta = 1 case was filtered above so theta - 1 != 0 here. + if constexpr (is_stan_scalar_v) { + partials<0>(ops_partials) = sum((n_arr + 1.0) * inv(theta_arr - 1.0)); + } else { + partials<0>(ops_partials) = (n_arr + 1.0) * inv(theta_arr - 1.0); } - 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); } + + return ops_partials.build(logP); } } // namespace math diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index 0032d0c7ec4..0737ae4ad79 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -3,25 +3,29 @@ #include #include -#include -#include -#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include #include -#include -#include -#include -#include +#include namespace stan { namespace math { /** \ingroup prob_dists * Returns the log CDF of the geometric distribution. Given containers of - * matching sizes, returns the log sum of probabilities. + * matching sizes, returns the log of the product of probabilities. * - * Delegates to the negative binomial log CDF with alpha = 1 and - * beta = theta / (1 - theta). + * 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 @@ -29,60 +33,66 @@ namespace math { * @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 +template * = nullptr> inline return_type_t geometric_lcdf(const T_n& n, const T_prob& theta) { - using T_n_ref = ref_type_t; - using T_prob_ref = ref_type_t; + using T_partials_return = partials_return_t; + using T_theta_ref = ref_type_t; 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_bounded(function, "Probability parameter", theta_arr, 0.0, 1.0); - check_consistent_sizes(function, "Outcome variable", n, - "Success probability parameter", theta); if (size_zero(n, theta)) { return 0.0; } - T_n_ref n_ref = n; - T_prob_ref theta_ref = theta; - check_bounded(function, "Success probability parameter", value_of(theta_ref), - 0.0, 1.0); + auto ops_partials = make_partials_propagator(theta_ref); - scalar_seq_view n_vec(n_ref); - for (int i = 0; i < stan::math::size(n); i++) { - if (n_vec.val(i) < 0) { - return negative_infinity(); - } + // log P(N <= n) = -inf for n < 0 + if (any(n_arr < 0)) { + return ops_partials.build(NEGATIVE_INFTY); } - // theta = 1 => log CDF is 0 for n >= 0 - scalar_seq_view 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; + // theta = 0 is degenerate: log P = -inf and the partials path divides + // by P_i = 0. Short-circuit. + if (any(theta_arr == 0.0)) { + return ops_partials.build(NEGATIVE_INFTY); } - if constexpr (is_stan_scalar_v) { - const auto beta = theta_ref / (1.0 - theta_ref); - return neg_binomial_lcdf(n_ref, 1, beta); - } else if constexpr (is_std_vector_v) { - std::vector> 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])); + // 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) { + // 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) { + partials<0>(ops_partials) = sum(elt_divide(dP_dtheta, P_i)); + } else { + partials<0>(ops_partials) = elt_divide(dP_dtheta, P_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); } + + return ops_partials.build(logP); } } // namespace math diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index df0a8f2a03e..40af9931022 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -3,15 +3,18 @@ #include #include +#include +#include +#include +#include +#include #include #include -#include +#include #include #include -#include -#include -#include -#include +#include +#include namespace stan { namespace math { @@ -20,8 +23,8 @@ namespace math { * Returns the log PMF of the geometric distribution. If containers * of matching sizes are supplied, returns the log sum of probabilities. * - * The geometric distribution is a special case of the negative - * binomial distribution with alpha = 1 and beta = theta / (1 - theta). + * The geometric distribution counts the number of failures before + * the first success: P(N = n | theta) = theta * (1 - theta)^n. * * @tparam T_n type of outcome variable * @tparam T_prob type of success probability parameter @@ -29,7 +32,7 @@ namespace math { * @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::domain_error if n is negative * @throw std::invalid_argument if container sizes mismatch */ @@ -37,55 +40,61 @@ template * = nullptr> inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { + using T_partials_return = partials_return_t; using T_n_ref = ref_type_t; - using T_prob_ref = ref_type_t; + using T_theta_ref = ref_type_t; static constexpr const char* function = "geometric_lpmf"; check_consistent_sizes(function, "Outcome variable", n, "Success probability parameter", theta); - if (size_zero(n, theta)) { - return 0.0; - } - T_n_ref n_ref = n; - T_prob_ref theta_ref = theta; + T_theta_ref theta_ref = theta; check_nonnegative(function, "Outcome variable", n_ref); check_bounded(function, "Success probability parameter", value_of(theta_ref), 0.0, 1.0); - // theta = 1 => deterministic: P(0) = 1, P(n>0) = 0 - // Cannot delegate since beta = theta / (1 - theta) diverges - scalar_seq_view n_vec(n_ref); - scalar_seq_view theta_vec(theta_ref); - const size_t max_sz = max_size(n_ref, theta_ref); - bool all_theta_one = true; - for (size_t i = 0; i < max_sz; i++) { - if (value_of(theta_vec[i]) == 1.0) { - if (n_vec[i] > 0) { - return negative_infinity(); - } - } else { - all_theta_one = false; - } + if (size_zero(n, theta)) { + return 0.0; } - if (all_theta_one) { + if constexpr (!include_summand::value) { return 0.0; } - // geometric(theta) = neg_binomial(1, theta / (1 - theta)) - if constexpr (is_stan_scalar_v) { - const auto beta = theta_ref / (1.0 - theta_ref); - return neg_binomial_lpmf(n_ref, 1, beta); - } else if constexpr (is_std_vector_v) { - std::vector> 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])); + auto ops_partials = make_partials_propagator(theta_ref); + + // Probability-zero event: theta = 1 with n > 0 + // (when theta = 1 success is certain, so only n = 0 has positive mass). + // A loop check here keeps both scalar and vector paths simple and + // avoids ambiguities in Eigen boolean array combination. + scalar_seq_view n_vec(n_ref); + scalar_seq_view 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); } - return neg_binomial_lpmf(n_ref, 1, beta); - } else { - const auto beta = elt_divide(theta_ref, subtract(1.0, theta_ref)); - return neg_binomial_lpmf(n_ref, 1, beta); } + + const auto& n_arr = as_value_column_array_or_scalar(n_ref); + const auto& theta_arr = as_value_column_array_or_scalar(theta_ref); + + // log P = log(theta) + n * log1m(theta). + // The select on n == 0 avoids 0 * log1m(1) = 0 * (-inf) = NaN when + // theta = 1; the n > 0 && theta = 1 case was already handled above. + const auto& log1m_theta = log1m(theta_arr); + const auto& failure_term = select(n_arr == 0, T_partials_return(0), + n_arr * log1m_theta); + T_partials_return logp = sum(log(theta_arr) + failure_term); + + if constexpr (is_autodiff_v) { + // d/dtheta log P = 1/theta - n/(1 - theta) = 1/theta + n/(theta - 1). + // For n = 0 the failure term contributes nothing; for n > 0 we have + // theta != 1 here, so theta - 1 != 0. + const auto& failure_grad = select(n_arr == 0, T_partials_return(0), + n_arr * inv(theta_arr - 1.0)); + partials<0>(ops_partials) = inv(theta_arr) + failure_grad; + } + + return ops_partials.build(logp); } template From 507a8f58325573bc6447b1a4237ef07579adb9ed Mon Sep 17 00:00:00 2001 From: Stan Jenkins Date: Fri, 24 Apr 2026 22:17:00 -0400 Subject: [PATCH 09/13] [Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1 --- stan/math/prim/prob/geometric_lpmf.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index 40af9931022..deb1bc77347 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -81,8 +81,8 @@ inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { // The select on n == 0 avoids 0 * log1m(1) = 0 * (-inf) = NaN when // theta = 1; the n > 0 && theta = 1 case was already handled above. const auto& log1m_theta = log1m(theta_arr); - const auto& failure_term = select(n_arr == 0, T_partials_return(0), - n_arr * log1m_theta); + const auto& failure_term + = select(n_arr == 0, T_partials_return(0), n_arr * log1m_theta); T_partials_return logp = sum(log(theta_arr) + failure_term); if constexpr (is_autodiff_v) { From 31dcbd6ed5703cfde07c51fa48f6d93517726313 Mon Sep 17 00:00:00 2001 From: yasumorishima Date: Sat, 25 Apr 2026 03:42:19 +0000 Subject: [PATCH 10/13] Tighten geometric distribution theta bound to (0, 1] Address CodeRabbit review on PR #3299: with check_bounded(theta, 0.0, 1.0) (inclusive), passing theta = 0 produced -inf logp but the autodiff path still evaluated inv(theta_arr), poisoning the partial with +inf. Replace the inclusive bound with check_positive_finite + check_less_or_ equal in lpmf/cdf/lcdf/lccdf to match the documented domain (0, 1] and the existing geometric_rng pattern. With the tightened bound the theta == 0 short-circuits in cdf/lcdf become unreachable; remove them. Add an EXPECT_THROW unit test covering theta = 0 across all four distribution functions to lock in the new behavior. --- stan/math/prim/prob/geometric_cdf.hpp | 10 +++------- stan/math/prim/prob/geometric_lccdf.hpp | 6 +++--- stan/math/prim/prob/geometric_lcdf.hpp | 10 +++------- stan/math/prim/prob/geometric_lpmf.hpp | 8 +++++--- test/unit/math/prim/prob/geometric_test.cpp | 11 +++++++++++ 5 files changed, 25 insertions(+), 20 deletions(-) diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index 17f38648d0f..692cd02145d 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -33,7 +33,7 @@ namespace math { * @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 geometric_cdf(const T_n& n, const T_prob& 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_bounded(function, "Probability parameter", theta_arr, 0.0, 1.0); + 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; @@ -61,11 +62,6 @@ inline return_type_t geometric_cdf(const T_n& n, const T_prob& theta) { return ops_partials.build(0.0); } - // theta = 0 is degenerate: P(N <= n) = 0 for any finite n. - // Avoid divide-by-zero in the partials path below. - if (any(theta_arr == 0.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), diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index 2532a59c2d4..a794e4c068c 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -30,7 +30,7 @@ namespace math { * @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::domain_error if theta is not in (0, 1] * @throw std::invalid_argument if container sizes mismatch */ template geometric_lccdf(const T_n& n, 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_bounded(function, "Probability parameter", theta_arr, 0.0, 1.0); + 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; @@ -75,7 +76,6 @@ inline return_type_t geometric_lccdf(const T_n& n, } // log P(N > n) = (n + 1) * log1m(theta) - // For theta = 0: log1m(0) = 0, lccdf = 0 (correct: certain failure). const auto& log1m_theta = log1m(theta_arr); T_partials_return logP = sum((n_arr + 1.0) * log1m_theta); diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index 0737ae4ad79..aec4ee62c11 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -33,7 +33,7 @@ namespace math { * @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 geometric_lcdf(const T_n& n, const T_prob& 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_bounded(function, "Probability parameter", theta_arr, 0.0, 1.0); + 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; @@ -61,11 +62,6 @@ inline return_type_t geometric_lcdf(const T_n& n, const T_prob& theta) { return ops_partials.build(NEGATIVE_INFTY); } - // theta = 0 is degenerate: log P = -inf and the partials path divides - // by P_i = 0. Short-circuit. - if (any(theta_arr == 0.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) diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index deb1bc77347..8be743097f4 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -32,7 +32,7 @@ namespace math { * @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::domain_error if n is negative * @throw std::invalid_argument if container sizes mismatch */ @@ -49,8 +49,10 @@ inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { T_n_ref n_ref = n; T_theta_ref theta_ref = theta; check_nonnegative(function, "Outcome variable", n_ref); - check_bounded(function, "Success probability parameter", value_of(theta_ref), - 0.0, 1.0); + check_positive_finite(function, "Success probability parameter", + value_of(theta_ref)); + check_less_or_equal(function, "Success probability parameter", + value_of(theta_ref), 1.0); if (size_zero(n, theta)) { return 0.0; diff --git a/test/unit/math/prim/prob/geometric_test.cpp b/test/unit/math/prim/prob/geometric_test.cpp index e564f78cfbf..081feaa7ff6 100644 --- a/test/unit/math/prim/prob/geometric_test.cpp +++ b/test/unit/math/prim/prob/geometric_test.cpp @@ -101,3 +101,14 @@ TEST(ProbDistributionsGeometric, rng_theta_one) { EXPECT_EQ(stan::math::geometric_rng(1.0, rng), 0); } } + +TEST(ProbDistributionsGeometric, theta_zero_throws) { + // theta = 0 is a degenerate parameter (no successes possible) and is + // outside the documented domain (0, 1]; all four distribution functions + // must throw rather than returning a poisoned partial. + EXPECT_THROW(stan::math::geometric_lpmf(0, 0.0), std::domain_error); + EXPECT_THROW(stan::math::geometric_lpmf(3, 0.0), std::domain_error); + EXPECT_THROW(stan::math::geometric_cdf(0, 0.0), std::domain_error); + EXPECT_THROW(stan::math::geometric_lcdf(0, 0.0), std::domain_error); + EXPECT_THROW(stan::math::geometric_lccdf(0, 0.0), std::domain_error); +} From 6337a62484e32479dc37b76d7b5b1337d57afbc7 Mon Sep 17 00:00:00 2001 From: Stan Jenkins Date: Fri, 24 Apr 2026 23:44:52 -0400 Subject: [PATCH 11/13] [Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1 --- stan/math/prim/prob/geometric_cdf.hpp | 1 - stan/math/prim/prob/geometric_lcdf.hpp | 1 - 2 files changed, 2 deletions(-) diff --git a/stan/math/prim/prob/geometric_cdf.hpp b/stan/math/prim/prob/geometric_cdf.hpp index 692cd02145d..7ac60b26388 100644 --- a/stan/math/prim/prob/geometric_cdf.hpp +++ b/stan/math/prim/prob/geometric_cdf.hpp @@ -62,7 +62,6 @@ inline return_type_t geometric_cdf(const T_n& n, const T_prob& theta) { 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 diff --git a/stan/math/prim/prob/geometric_lcdf.hpp b/stan/math/prim/prob/geometric_lcdf.hpp index aec4ee62c11..84870f31ec5 100644 --- a/stan/math/prim/prob/geometric_lcdf.hpp +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -62,7 +62,6 @@ inline return_type_t geometric_lcdf(const T_n& n, const T_prob& theta) { 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 From b37a95835910a828162d0be446bb1080e8df68e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yasunori=20Morishima=EF=BC=88=E7=9B=9B=E5=B3=B6=E5=BA=B7?= =?UTF-8?q?=E5=BE=B3=EF=BC=89?= Date: Sat, 25 Apr 2026 04:50:24 +0000 Subject: [PATCH 12/13] Fix lccdf vectorized contracts and lpmf scalar partials geometric_lccdf: - Drop the any(n_arr < 0) early return that silently dropped contributions from positive n_i in mixed-sign vectors. P(N > n_i) = 1 for n_i < 0 is the multiplicative identity (log 0), not an absorbing element, so the term has to remain in the sum. Per-element select(n < 0, 0, ...) keeps both value and partials paths correct. - Replace the any(theta == 1.0) blanket guard with a scalar_seq_view per-element loop, so theta=1 only forces -inf when paired with n_i >= 0 at the same index. Previously theta=[0.5, 1.0], n=[2, -1] returned -inf instead of 3*log(0.5). geometric_lpmf: - Add the is_stan_scalar_v guard around partials<0> assignment so a scalar theta paired with a vector n collapses to a scalar partial via sum(...), matching cdf/lcdf/lccdf. Tests: lccdf_vectorized_mixed_sign and lccdf_vectorized_theta_one_alignment lock in both contract fixes. --- stan/math/prim/prob/geometric_lccdf.hpp | 49 +++++++++++++-------- stan/math/prim/prob/geometric_lpmf.hpp | 7 ++- test/unit/math/prim/prob/geometric_test.cpp | 20 +++++++++ 3 files changed, 56 insertions(+), 20 deletions(-) diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index a794e4c068c..61ac05a66aa 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -8,6 +8,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -39,12 +42,14 @@ template geometric_lccdf(const T_n& n, const T_prob& theta) { using T_partials_return = partials_return_t; + using T_n_ref = ref_type_t; using T_theta_ref = ref_type_t; 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); + 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); @@ -55,12 +60,6 @@ inline return_type_t geometric_lccdf(const T_n& n, auto ops_partials = make_partials_propagator(theta_ref); - // log P(N > n) = 0 (i.e. P = 1) when n < 0, matching the existing - // implementation that short-circuits on the first negative element. - if (any(n_arr < 0)) { - return ops_partials.build(0.0); - } - // 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.) @@ -68,25 +67,37 @@ inline return_type_t geometric_lccdf(const T_n& n, return ops_partials.build(NEGATIVE_INFTY); } - // theta = 1 means certain success, so P(N > n) = 0 for n >= 0 and the - // log is -inf. The partials path divides by (theta - 1) = 0, so we - // short-circuit. - if (any(theta_arr == 1.0)) { - 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 n_vec(n_ref); + scalar_seq_view 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); + } } - // log P(N > n) = (n + 1) * log1m(theta) + // 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); - T_partials_return logP = sum((n_arr + 1.0) * log1m_theta); + 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) { - // d/dtheta (n + 1) * log1m(theta) = -(n + 1) / (1 - theta) - // = (n + 1) / (theta - 1) - // theta = 1 case was filtered above so theta - 1 != 0 here. + // 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) { - partials<0>(ops_partials) = sum((n_arr + 1.0) * inv(theta_arr - 1.0)); + partials<0>(ops_partials) = sum(dterm); } else { - partials<0>(ops_partials) = (n_arr + 1.0) * inv(theta_arr - 1.0); + partials<0>(ops_partials) = dterm; } } diff --git a/stan/math/prim/prob/geometric_lpmf.hpp b/stan/math/prim/prob/geometric_lpmf.hpp index 8be743097f4..7649a4bb646 100644 --- a/stan/math/prim/prob/geometric_lpmf.hpp +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -93,7 +94,11 @@ inline return_type_t geometric_lpmf(const T_n& n, const T_prob& theta) { // theta != 1 here, so theta - 1 != 0. const auto& failure_grad = select(n_arr == 0, T_partials_return(0), n_arr * inv(theta_arr - 1.0)); - partials<0>(ops_partials) = inv(theta_arr) + failure_grad; + if constexpr (is_stan_scalar_v) { + partials<0>(ops_partials) = sum(inv(theta_arr) + failure_grad); + } else { + partials<0>(ops_partials) = inv(theta_arr) + failure_grad; + } } return ops_partials.build(logp); diff --git a/test/unit/math/prim/prob/geometric_test.cpp b/test/unit/math/prim/prob/geometric_test.cpp index 081feaa7ff6..df6ddc3ff8d 100644 --- a/test/unit/math/prim/prob/geometric_test.cpp +++ b/test/unit/math/prim/prob/geometric_test.cpp @@ -112,3 +112,23 @@ TEST(ProbDistributionsGeometric, theta_zero_throws) { EXPECT_THROW(stan::math::geometric_lcdf(0, 0.0), std::domain_error); EXPECT_THROW(stan::math::geometric_lccdf(0, 0.0), std::domain_error); } + +TEST(ProbDistributionsGeometric, lccdf_vectorized_mixed_sign) { + // n = [-1, 2], theta = 0.5: P(N > -1) = 1 (contributes log 1 = 0), + // P(N > 2) = 0.5^3 = 0.125 (contributes 3*log 0.5). + // Expected: 3 * log(0.5) ~= -2.0794. + std::vector n = {-1, 2}; + EXPECT_NEAR(stan::math::geometric_lccdf(n, 0.5), 3.0 * std::log(0.5), 1e-10); +} + +TEST(ProbDistributionsGeometric, lccdf_vectorized_theta_one_alignment) { + // theta_i = 1 must only zero out the result when paired with n_i >= 0. + // n = [2, -1], theta = [0.5, 1.0]: + // i=0: theta=0.5, n= 2 -> 3 * log(0.5) + // i=1: theta=1.0, n=-1 -> 0 (n below support) + // Expected: 3 * log(0.5). + std::vector n = {2, -1}; + std::vector theta = {0.5, 1.0}; + EXPECT_NEAR(stan::math::geometric_lccdf(n, theta), 3.0 * std::log(0.5), + 1e-10); +} From 97cb3ef197f4ac13438ee95f6e9a739e507973fc Mon Sep 17 00:00:00 2001 From: Stan Jenkins Date: Sat, 25 Apr 2026 00:52:00 -0400 Subject: [PATCH 13/13] [Jenkins] auto-formatting by clang-format version 10.0.0-4ubuntu1 --- stan/math/prim/prob/geometric_lccdf.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/stan/math/prim/prob/geometric_lccdf.hpp b/stan/math/prim/prob/geometric_lccdf.hpp index 61ac05a66aa..0e565e6d935 100644 --- a/stan/math/prim/prob/geometric_lccdf.hpp +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -85,8 +85,8 @@ inline return_type_t geometric_lccdf(const T_n& n, // 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); + 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) {