Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add beta_neg_binomial_lccdf #3114

Merged
merged 8 commits into from
Oct 27, 2024
1 change: 1 addition & 0 deletions stan/math/prim/prob.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <stan/math/prim/prob/beta_lccdf.hpp>
#include <stan/math/prim/prob/beta_lcdf.hpp>
#include <stan/math/prim/prob/beta_lpdf.hpp>
#include <stan/math/prim/prob/beta_neg_binomial_lccdf.hpp>
#include <stan/math/prim/prob/beta_neg_binomial_lpmf.hpp>
#include <stan/math/prim/prob/beta_proportion_ccdf_log.hpp>
#include <stan/math/prim/prob/beta_proportion_cdf_log.hpp>
Expand Down
151 changes: 151 additions & 0 deletions stan/math/prim/prob/beta_neg_binomial_lccdf.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
#ifndef STAN_MATH_PRIM_PROB_BETA_NEG_BINOMIAL_LCCDF_HPP
#define STAN_MATH_PRIM_PROB_BETA_NEG_BINOMIAL_LCCDF_HPP

#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/digamma.hpp>
#include <stan/math/prim/fun/hypergeometric_3F2.hpp>
#include <stan/math/prim/fun/grad_F32.hpp>
#include <stan/math/prim/fun/lbeta.hpp>
#include <stan/math/prim/fun/lgamma.hpp>
#include <stan/math/prim/fun/max_size.hpp>
#include <stan/math/prim/fun/scalar_seq_view.hpp>
#include <stan/math/prim/fun/size.hpp>
#include <stan/math/prim/fun/size_zero.hpp>
#include <stan/math/prim/functor/partials_propagator.hpp>
#include <cmath>

namespace stan {
namespace math {

/** \ingroup prob_dists
* Returns the log CCDF of the Beta-Negative Binomial distribution with given
* number of successes, prior success, and prior failure parameters.
* Given containers of matching sizes, returns the log sum of probabilities.
*
* @tparam T_n type of failure parameter
* @tparam T_r type of number of successes parameter
* @tparam T_alpha type of prior success parameter
* @tparam T_beta type of prior failure parameter
*
* @param n failure parameter
* @param r Number of successes parameter
* @param alpha prior success parameter
* @param beta prior failure parameter
* @return log probability or log sum of probabilities
* @throw std::domain_error if r, alpha, or beta fails to be positive
* @throw std::invalid_argument if container sizes mismatch
*/
template <typename T_n, typename T_r, typename T_alpha, typename T_beta>
inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lccdf(
const T_n& n, const T_r& r, const T_alpha& alpha, const T_beta& beta) {
using std::exp;
using std::log;
using T_partials_return = partials_return_t<T_n, T_r, T_alpha, T_beta>;
using T_r_ref = ref_type_t<T_r>;
using T_alpha_ref = ref_type_t<T_alpha>;
using T_beta_ref = ref_type_t<T_beta>;
static constexpr const char* function = "beta_neg_binomial_lccdf";
check_consistent_sizes(
function, "Failures variable", n, "Number of successes parameter", r,
"Prior success parameter", alpha, "Prior failure parameter", beta);
if (size_zero(n, r, alpha, beta)) {
return 0;
}

T_r_ref r_ref = r;
T_alpha_ref alpha_ref = alpha;
T_beta_ref beta_ref = beta;
check_positive_finite(function, "Number of successes parameter", r_ref);
check_positive_finite(function, "Prior success parameter", alpha_ref);
check_positive_finite(function, "Prior failure parameter", beta_ref);

T_partials_return P(0.0);
lingium marked this conversation as resolved.
Show resolved Hide resolved
auto ops_partials = make_partials_propagator(r_ref, alpha_ref, beta_ref);

scalar_seq_view<T_n> n_vec(n);
scalar_seq_view<T_r_ref> r_vec(r_ref);
scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
scalar_seq_view<T_beta_ref> beta_vec(beta_ref);
size_t size_n = stan::math::size(n);
size_t max_size_seq_view = max_size(n, r, alpha, beta);

// Explicit return for extreme values
// The gradients are technically ill-defined, but treated as zero
for (size_t i = 0; i < size_n; i++) {
if (n_vec.val(i) < 0) {
return ops_partials.build(0.0);
}
}
lingium marked this conversation as resolved.
Show resolved Hide resolved

for (size_t i = 0; i < max_size_seq_view; i++) {
// Explicit return for extreme values
// The gradients are technically ill-defined, but treated as zero
if (n_vec.val(i) == std::numeric_limits<int>::max()) {
return ops_partials.build(negative_infinity());
}
lingium marked this conversation as resolved.
Show resolved Hide resolved
T_partials_return n_dbl = n_vec.val(i);
T_partials_return r_dbl = r_vec.val(i);
T_partials_return alpha_dbl = alpha_vec.val(i);
T_partials_return beta_dbl = beta_vec.val(i);
T_partials_return b_plus_n = beta_dbl + n_dbl;
T_partials_return r_plus_n = r_dbl + n_dbl;
T_partials_return a_plus_r = alpha_dbl + r_dbl;
T_partials_return one = 1;
T_partials_return precision = 1e-8; // default -6, set -8 to pass all tests
lingium marked this conversation as resolved.
Show resolved Hide resolved

T_partials_return F
= hypergeometric_3F2({one, b_plus_n + 1, r_plus_n + 1},
{n_dbl + 2, a_plus_r + b_plus_n + 1}, one);
T_partials_return C = lgamma(r_plus_n + 1) + lbeta(a_plus_r, b_plus_n + 1)
- lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl)
- lgamma(n_dbl + 2);
T_partials_return ccdf = exp(C) * F;
T_partials_return P_i = log(ccdf);
P += P_i;

if (!is_constant_all<T_r, T_alpha, T_beta>::value) {
lingium marked this conversation as resolved.
Show resolved Hide resolved
T_partials_return digamma_n_r_alpha_beta
= digamma(a_plus_r + b_plus_n + 1);
T_partials_return dF[6];
grad_F32(dF, one, b_plus_n + 1, r_plus_n + 1, n_dbl + 2,
a_plus_r + b_plus_n + 1, one, precision, 1e5);
SteveBronder marked this conversation as resolved.
Show resolved Hide resolved

if constexpr (!is_constant<T_r>::value || !is_constant<T_alpha>::value) {
T_partials_return digamma_r_alpha = digamma(a_plus_r);
if constexpr (!is_constant_all<T_r>::value) {
partials<0>(ops_partials)[i]
+= digamma(r_plus_n + 1)
+ (digamma_r_alpha - digamma_n_r_alpha_beta)
+ (dF[2] + dF[4]) / F - digamma(r_dbl);
}
if constexpr (!is_constant_all<T_alpha>::value) {
partials<1>(ops_partials)[i] += digamma_r_alpha
- digamma_n_r_alpha_beta + dF[4] / F
- digamma(alpha_dbl);
}
}

if constexpr (!is_constant<T_alpha>::value
|| !is_constant<T_beta>::value) {
T_partials_return digamma_alpha_beta = digamma(alpha_dbl + beta_dbl);
if constexpr (!is_constant<T_alpha>::value) {
partials<1>(ops_partials)[i] += digamma_alpha_beta;
}
if constexpr (!is_constant<T_beta>::value) {
partials<2>(ops_partials)[i]
+= digamma(b_plus_n + 1) - digamma_n_r_alpha_beta
+ (dF[1] + dF[4]) / F
- (digamma(beta_dbl) - digamma_alpha_beta);
}
}
}
}

return ops_partials.build(P);
}

} // namespace math
} // namespace stan
#endif
96 changes: 96 additions & 0 deletions test/prob/beta_neg_binomial/beta_neg_binomial_ccdf_log_test.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
// Arguments: Ints, Doubles, Doubles, Doubles
#include <stan/math/prim/prob/beta_neg_binomial_lccdf.hpp>
#include <stan/math/prim/fun/lbeta.hpp>
#include <stan/math/prim/fun/lgamma.hpp>

using stan::math::var;
using std::numeric_limits;
using std::vector;

class AgradCcdfLogBetaNegBinomial : public AgradCcdfLogTest {
public:
void valid_values(vector<vector<double>>& parameters,
vector<double>& ccdf_log) {
vector<double> param(4);

param[0] = 10; // n
param[1] = 5.5; // r
param[2] = 2.5; // alpha
param[3] = 0.5; // beta
Comment on lines +16 to +19
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add some test for more extreme values? Mostly so we can make sure we have good precision at the bounds of the support

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will add more in the next commit

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Didn't add too many tests. Cuz if the values are extreme, grad_F32 would be too complex. We have to increase max allowed iterations to pass tests, then the tests might run for years lol

Copy link
Collaborator Author

@lingium lingium Oct 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

though I do test on local for all parameters = 1e-3 or 1e3, increase max_steps to 1e-8 can pass tests

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thats fine. Let's leave it at 1e-8 because we don't really know where the algorithm / user is going to end up in the parameter space.

parameters.push_back(param);
ccdf_log.push_back(std::log(1.0 - 0.967906252841089)); // expected ccdf_log
}

void invalid_values(vector<size_t>& index, vector<double>& value) {
// n

// r
index.push_back(1U);
value.push_back(0.0);

index.push_back(1U);
value.push_back(-1.0);

index.push_back(1U);
value.push_back(std::numeric_limits<double>::infinity());

// alpha
index.push_back(2U);
value.push_back(0.0);

index.push_back(2U);
value.push_back(-1.0);

index.push_back(2U);
value.push_back(std::numeric_limits<double>::infinity());

// beta
index.push_back(3U);
value.push_back(0.0);

index.push_back(3U);
value.push_back(-1.0);

index.push_back(3U);
value.push_back(std::numeric_limits<double>::infinity());
}

// BOUND INCLUDED IN ORDER FOR TEST TO PASS WITH CURRENT FRAMEWORK
bool has_lower_bound() { return false; }

bool has_upper_bound() { return false; }

template <typename T_n, typename T_r, typename T_size1, typename T_size2,
typename T4, typename T5>
stan::return_type_t<T_r, T_size1, T_size2> ccdf_log(const T_n& n,
const T_r& r,
const T_size1& alpha,
const T_size2& beta,
const T4&, const T5&) {
return stan::math::beta_neg_binomial_lccdf(n, r, alpha, beta);
}

template <typename T_n, typename T_r, typename T_size1, typename T_size2,
typename T4, typename T5>
stan::return_type_t<T_r, T_size1, T_size2> ccdf_log_function(
const T_n& n, const T_r& r, const T_size1& alpha, const T_size2& beta,
const T4&, const T5&) {
using stan::math::lbeta;
using stan::math::lgamma;
using stan::math::log1m;
using stan::math::log_sum_exp;
using std::vector;

vector<stan::return_type_t<T_r, T_size1, T_size2>> lpmf_values;

for (int i = 0; i <= n; i++) {
auto lpmf = lbeta(i + r, alpha + beta) - lbeta(r, alpha)
+ lgamma(i + beta) - lgamma(i + 1) - lgamma(beta);
lpmf_values.push_back(lpmf);
}

auto log_cdf = log_sum_exp(lpmf_values);

return log1m(exp(log_cdf));
}
};