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 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] [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&,