diff --git a/stan/math/prim/prob.hpp b/stan/math/prim/prob.hpp index 308b084a165..33c39a183b8 100644 --- a/stan/math/prim/prob.hpp +++ b/stan/math/prim/prob.hpp @@ -113,6 +113,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..7ac60b26388 --- /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 +#include +#include +#include + +namespace stan { +namespace math { + +/** \ingroup prob_dists + * Returns the CDF of the geometric distribution. Given containers of + * matching sizes, returns the product of probabilities. + * + * The geometric distribution counts the number of failures before + * the first success: P(N <= n | theta) = 1 - (1 - theta)^(n + 1). + * + * @tparam T_n type of outcome variable + * @tparam T_prob type of success probability parameter + * + * @param n outcome variable (number of failures before first success) + * @param theta success probability parameter + * @return probability or product of probabilities + * @throw std::domain_error if theta is not in (0, 1] + * @throw std::invalid_argument if container sizes mismatch + */ +template * = nullptr> +inline return_type_t geometric_cdf(const T_n& n, const T_prob& theta) { + 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_positive_finite(function, "Probability parameter", theta_arr); + check_less_or_equal(function, "Probability parameter", theta_arr, 1.0); + + if (size_zero(n, theta)) { + return 1.0; + } + + auto ops_partials = make_partials_propagator(theta_ref); + + // P(N <= n) = 0 for n < 0 + if (any(n_arr < 0)) { + return ops_partials.build(0.0); + } + + // P_i = 1 - (1 - theta)^(n + 1) = -expm1((n + 1) * log1m(theta)) + // For theta = 1: log1m(1) = -inf, (n+1)*-inf = -inf (n >= 0), + // expm1(-inf) = -1, so P_i = 1 (correct: certain success means + // N <= n always for n >= 0). + const auto& log1m_theta = log1m(theta_arr); + const auto& P_i = -expm1((n_arr + 1.0) * log1m_theta); + const T_partials_return P = prod(P_i); + + if constexpr (is_autodiff_v) { + // 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 ops_partials.build(P); +} + +} // 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..0e565e6d935 --- /dev/null +++ b/stan/math/prim/prob/geometric_lccdf.hpp @@ -0,0 +1,109 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LCCDF_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_LCCDF_HPP + +#include +#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 of the product of complementary + * probabilities. + * + * log P(N > n | theta) = log((1 - theta)^(n + 1)) = (n + 1) * log1m(theta). + * + * @tparam T_n type of outcome variable + * @tparam T_prob type of success probability parameter + * + * @param n outcome variable (number of failures before first success) + * @param theta success probability parameter + * @return log complementary probability or log product of complements + * @throw std::domain_error if theta is not in (0, 1] + * @throw std::invalid_argument if container sizes mismatch + */ +template * = nullptr> +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_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_ref); + const auto& theta_arr = as_value_column_array_or_scalar(theta_ref); + check_positive_finite(function, "Probability parameter", theta_arr); + check_less_or_equal(function, "Probability parameter", theta_arr, 1.0); + + if (size_zero(n, theta)) { + return 0.0; + } + + auto ops_partials = make_partials_propagator(theta_ref); + + // n at INT_MAX: P(N > n) underflows to 0, lccdf = -inf. + // (The autodiff test framework probes the upper bound at INT_MAX, + // mirroring the early return used in neg_binomial_lccdf.) + if (any(n_arr == std::numeric_limits::max())) { + return ops_partials.build(NEGATIVE_INFTY); + } + + // theta = 1 with any n_i >= 0 makes (n_i + 1) * log1m(1) = -inf and + // the partials path divides by (theta - 1) = 0. Per-element loop here + // correctly pairs theta_i with n_i (e.g. theta=[0.5, 1.0], n=[2, -1] + // returns 3*log(0.5), not -inf). For theta_i = 1 paired with n_i < 0 + // the per-element select below contributes 0 to both value and partials. + scalar_seq_view 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); + } + } + + // Per-element: n_i < 0 contributes log(1) = 0 (since P(N > n_i) = 1 + // for any n_i below the support), so they are no-ops in the sum. + // n_i >= 0 contributes (n_i + 1) * log1m(theta). + const auto& log1m_theta = log1m(theta_arr); + const auto& term + = select(n_arr < 0, T_partials_return(0), (n_arr + 1.0) * log1m_theta); + T_partials_return logP = sum(term); + + if constexpr (is_autodiff_v) { + // 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(dterm); + } else { + partials<0>(ops_partials) = dterm; + } + } + + return ops_partials.build(logP); +} + +} // 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..84870f31ec5 --- /dev/null +++ b/stan/math/prim/prob/geometric_lcdf.hpp @@ -0,0 +1,95 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_LCDF_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_LCDF_HPP + +#include +#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 of the product of probabilities. + * + * log P(N <= n | theta) = log(1 - (1 - theta)^(n + 1)) + * = log1m_exp((n + 1) * log1m(theta)) + * + * @tparam T_n type of outcome variable + * @tparam T_prob type of success probability parameter + * + * @param n outcome variable (number of failures before first success) + * @param theta success probability parameter + * @return log probability or log sum of probabilities + * @throw std::domain_error if theta is not in (0, 1] + * @throw std::invalid_argument if container sizes mismatch + */ +template * = nullptr> +inline return_type_t geometric_lcdf(const T_n& n, const T_prob& theta) { + 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_positive_finite(function, "Probability parameter", theta_arr); + check_less_or_equal(function, "Probability parameter", theta_arr, 1.0); + + if (size_zero(n, theta)) { + return 0.0; + } + + auto ops_partials = make_partials_propagator(theta_ref); + + // log P(N <= n) = -inf for n < 0 + if (any(n_arr < 0)) { + return ops_partials.build(NEGATIVE_INFTY); + } + + // log_q = log((1 - theta)^(n + 1)) = (n + 1) * log1m(theta) + // log P_i = log(1 - q_i) = log1m_exp(log_q_i) + // For theta = 1: log_q = -inf, log1m_exp(-inf) = log(1 - 0) = 0 + // (correct: P = 1 when success is certain). + const auto& log1m_theta = log1m(theta_arr); + const auto& log_q = (n_arr + 1.0) * log1m_theta; + const auto& log_P_i = log1m_exp(log_q); + T_partials_return logP = sum(log_P_i); + + if constexpr (is_autodiff_v) { + // 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 ops_partials.build(logP); +} + +} // 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..7649a4bb646 --- /dev/null +++ b/stan/math/prim/prob/geometric_lpmf.hpp @@ -0,0 +1,114 @@ +#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 +#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 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 + * + * @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 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_lpmf"; + check_consistent_sizes(function, "Outcome variable", n, + "Success probability parameter", theta); + T_n_ref n_ref = n; + T_theta_ref theta_ref = theta; + check_nonnegative(function, "Outcome variable", n_ref); + 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; + } + if constexpr (!include_summand::value) { + return 0.0; + } + + 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); + } + } + + 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)); + 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); +} + +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..fa7619a86c6 --- /dev/null +++ b/stan/math/prim/prob/geometric_rng.hpp @@ -0,0 +1,63 @@ +#ifndef STAN_MATH_PRIM_PROB_GEOMETRIC_RNG_HPP +#define STAN_MATH_PRIM_PROB_GEOMETRIC_RNG_HPP + +#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. + * + * 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 + * + * @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); + + if constexpr (is_stan_scalar_v) { + if (theta_ref == 1.0) { + return 0; + } + 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()); + for (size_t i = 0; i < theta_arr.size(); i++) { + if (theta_arr[i] == 1.0) { + result[i] = 0; + } else { + double beta_i = theta_arr[i] / (1.0 - theta_arr[i]); + result[i] = neg_binomial_rng(1, beta_i, rng); + } + } + return result; + } +} + +} // namespace math +} // namespace stan +#endif 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 new file mode 100644 index 00000000000..cfd36f66c4d --- /dev/null +++ b/test/prob/geometric/geometric_cdf_test.hpp @@ -0,0 +1,62 @@ +// 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); + } + + 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/prob/geometric/geometric_test.hpp b/test/prob/geometric/geometric_test.hpp new file mode 100644 index 00000000000..8acf08fd22f --- /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); + } +}; 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..df6ddc3ff8d --- /dev/null +++ b/test/unit/math/prim/prob/geometric_test.cpp @@ -0,0 +1,134 @@ +#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); + } +} + +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); +} + +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); +}