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_lcdf #3116

Merged
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions stan/math/prim/prob/beta_neg_binomial_lccdf.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -34,15 +34,15 @@ namespace math {
* @param alpha prior success parameter
* @param beta prior failure parameter
* @param precision precision for `grad_F32`, default \f$10^{-8}\f$
* @param max_steps max iteration allowed for `grad_F32`, default \f$10^{-8}\f$
* @param max_steps max iteration allowed for `grad_F32`, default \f$10^{8}\f$
* @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,
const double precision = 1e-8, const int max_steps = 1e6) {
const double precision = 1e-8, const int max_steps = 1e8) {
static constexpr const char* function = "beta_neg_binomial_lccdf";
check_consistent_sizes(
function, "Failures variable", n, "Number of successes parameter", r,
Expand Down Expand Up @@ -100,7 +100,7 @@ inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lccdf(
std::initializer_list<b_t>{n_dbl + 2.0, a_plus_r + b_plus_n + 1.0},
1.0);
auto C = lgamma(r_plus_n + 1.0) + lbeta(a_plus_r, b_plus_n + 1.0)
- lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl) - lgamma(n_dbl + 2);
- lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl) - lgamma(n_dbl + 2.0);
log_ccdf += C + stan::math::log(F);

if constexpr (!is_constant_all<T_r, T_alpha, T_beta>::value) {
Expand Down
152 changes: 152 additions & 0 deletions stan/math/prim/prob/beta_neg_binomial_lcdf.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
#ifndef STAN_MATH_PRIM_PROB_BETA_NEG_BINOMIAL_LCDF_HPP
#define STAN_MATH_PRIM_PROB_BETA_NEG_BINOMIAL_LCDF_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 CDF 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
* @param precision precision for `grad_F32`, default \f$10^{-8}\f$
* @param max_steps max iteration allowed for `grad_F32`, default \f$10^{8}\f$
* @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_lcdf(
const T_n& n, const T_r& r, const T_alpha& alpha, const T_beta& beta,
const double precision = 1e-8, const int max_steps = 1e8) {
static constexpr const char* function = "beta_neg_binomial_lcdf";
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;
}

using T_r_ref = ref_type_t<T_r>;
T_r_ref r_ref = r;
using T_alpha_ref = ref_type_t<T_alpha>;
T_alpha_ref alpha_ref = alpha;
using T_beta_ref = ref_type_t<T_beta>;
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);

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);
int 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 (int i = 0; i < size_n; i++) {
if (n_vec.val(i) < 0) {
return negative_infinity();
}
}

using T_partials_return = partials_return_t<T_n, T_r, T_alpha, T_beta>;
T_partials_return log_cdf(0.0);
auto ops_partials = make_partials_propagator(r_ref, alpha_ref, beta_ref);
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 0.0;
}
auto n_dbl = n_vec.val(i);
auto r_dbl = r_vec.val(i);
auto alpha_dbl = alpha_vec.val(i);
auto beta_dbl = beta_vec.val(i);
auto b_plus_n = beta_dbl + n_dbl;
auto r_plus_n = r_dbl + n_dbl;
auto a_plus_r = alpha_dbl + r_dbl;
using a_t = return_type_t<decltype(b_plus_n), decltype(r_plus_n)>;
using b_t = return_type_t<decltype(n_dbl), decltype(a_plus_r),
decltype(b_plus_n)>;
auto F = hypergeometric_3F2(
std::initializer_list<a_t>{1.0, b_plus_n + 1.0, r_plus_n + 1.0},
std::initializer_list<b_t>{n_dbl + 2.0, a_plus_r + b_plus_n + 1.0},
1.0);
auto C = lgamma(r_plus_n + 1.0) + lbeta(a_plus_r, b_plus_n + 1.0)
- lgamma(r_dbl) - lbeta(alpha_dbl, beta_dbl) - lgamma(n_dbl + 2.0);
auto ccdf = stan::math::exp(C) * F;
log_cdf += log1m(ccdf);

if constexpr (!is_constant_all<T_r, T_alpha, T_beta>::value) {
auto chain_rule_term = -ccdf / (1.0 - ccdf);
auto digamma_n_r_alpha_beta = digamma(a_plus_r + b_plus_n + 1.0);
T_partials_return dF[6];
grad_F32<false, !is_constant<T_beta>::value, !is_constant_all<T_r>::value,
false, true, false>(dF, 1.0, b_plus_n + 1.0, r_plus_n + 1.0,
n_dbl + 2.0, a_plus_r + b_plus_n + 1.0, 1.0,
precision, max_steps);

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

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

return ops_partials.build(log_cdf);
}

} // namespace math
} // namespace stan
#endif
92 changes: 92 additions & 0 deletions test/prob/beta_neg_binomial/beta_neg_binomial_cdf_log_test.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
// Arguments: Ints, Doubles, Doubles, Doubles
#include <stan/math/prim/prob/beta_neg_binomial_lcdf.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 AgradCdfLogBetaNegBinomial : public AgradCdfLogTest {
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
parameters.push_back(param);
ccdf_log.push_back(std::log(0.967906252841089)); // expected cdf_log
lingium marked this conversation as resolved.
Show resolved Hide resolved
}

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> cdf_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_lcdf(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> cdf_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::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);
}

return log_sum_exp(lpmf_values);
}
};