Skip to content

Commit

Permalink
Some updates for NegativeBinomial (#164)
Browse files Browse the repository at this point in the history 
* Uncomment the check_discrete_distribution tests.
* Expand the documentation--explain the meaning of `r` and `p`.
* Note that `r` can be real.

Closes gh-115.

Co-authored-by: Michael Ma <michael60612@gmail.com>
  • Loading branch information
@WarrenWeckesser @boxtown
WarrenWeckesser and boxtown committed Aug 29, 2022
1 parent e00a6e5 commit 184367c
Showing 1 changed file with 65 additions and 36 deletions.
101 changes: 65 additions & 36 deletions src/distribution/negative_binomial.rs
View file
Expand Up @@ -6,8 +6,22 @@ use rand::Rng;
use std::f64;

/// Implements the
/// [NegativeBinomial](http://en.wikipedia.org/wiki/Negative_binomial_distribution)
/// distribution
/// [negative binomial](http://en.wikipedia.org/wiki/Negative_binomial_distribution)
/// distribution.
///
/// *Please note carefully the meaning of the parameters.* As noted in the
/// wikipedia article, there are several different commonly used conventions
/// for the parameters of the negative binomial distribution.
///
/// The negative binomial distribution is a discrete distribution with two
/// parameters, `r` and `p`. When `r` is an integer, the negative binomial
/// distribution can be interpreted as the distribution of the number of
/// failures in a sequence of Bernoulli trials that continue until `r`
/// successes occur. `p` is the probability of success in a single Bernoulli
/// trial.
///
/// `NegativeBinomial` accepts non-integer values for `r`. This is a
/// generalization of the more common case where `r` is an integer.
///
/// # Examples
///
Expand All @@ -28,8 +42,11 @@ pub struct NegativeBinomial {
}

impl NegativeBinomial {
/// Constructs a new negative binomial distribution
/// with a given `p` probability of the number of successes `r`
/// Constructs a new negative binomial distribution with parameters `r`
/// and `p`. When `r` is an integer, the negative binomial distribution
/// can be interpreted as the distribution of the number of failures in
/// a sequence of Bernoulli trials that continue until `r` successes occur.
/// `p` is the probability of success in a single Bernoulli trial.
///
/// # Errors
///
Expand All @@ -55,8 +72,9 @@ impl NegativeBinomial {
}
}

/// Returns the probability of success `p` of
/// the negative binomial distribution.
/// Returns the probability of success `p` of a single
/// Bernoulli trial associated with the negative binomial
/// distribution.
///
/// # Examples
///
Expand All @@ -71,7 +89,7 @@ impl NegativeBinomial {
}

/// Returns the number `r` of success of this negative
/// binomial distribution
/// binomial distribution.
///
/// # Examples
///
Expand All @@ -95,21 +113,15 @@ impl ::rand::distributions::Distribution<u64> for NegativeBinomial {

impl DiscreteCDF<u64, f64> for NegativeBinomial {
/// Calculates the cumulative distribution function for the
/// negative binomial distribution at `x`
///
/// Note that due to extending the distribution to the reals
/// (allowing positive real values for `r`), while still technically
/// a discrete distribution the CDF behaves more like that of a
/// continuous distribution rather than a discrete distribution
/// (i.e. a smooth graph rather than a step-ladder)
/// negative binomial distribution at `x`.
///
/// # Formula
///
/// ```ignore
/// I_(p)(r, x+1)
/// ```
///
/// where `I_(x)(a, b)` is the regularized incomplete beta function
/// where `I_(x)(a, b)` is the regularized incomplete beta function.
fn cdf(&self, x: u64) -> f64 {
beta::beta_reg(self.r, x as f64 + 1.0, self.p)
}
Expand Down Expand Up @@ -138,7 +150,7 @@ impl DiscreteCDF<u64, f64> for NegativeBinomial {
impl Min<u64> for NegativeBinomial {
/// Returns the minimum value in the domain of the
/// negative binomial distribution representable by a 64-bit
/// integer
/// integer.
///
/// # Formula
///
Expand All @@ -153,7 +165,7 @@ impl Min<u64> for NegativeBinomial {
impl Max<u64> for NegativeBinomial {
/// Returns the maximum value in the domain of the
/// negative binomial distribution representable by a 64-bit
/// integer
/// integer.
///
/// # Formula
///
Expand All @@ -166,7 +178,7 @@ impl Max<u64> for NegativeBinomial {
}

impl DiscreteDistribution<f64> for NegativeBinomial {
/// Returns the mean of the negative binomial distribution
/// Returns the mean of the negative binomial distribution.
///
/// # Formula
///
Expand All @@ -176,7 +188,7 @@ impl DiscreteDistribution<f64> for NegativeBinomial {
fn mean(&self) -> Option<f64> {
Some(self.r * (1.0 - self.p) / self.p)
}
/// Returns the variance of the negative binomial distribution
/// Returns the variance of the negative binomial distribution.
///
/// # Formula
///
Expand All @@ -186,7 +198,7 @@ impl DiscreteDistribution<f64> for NegativeBinomial {
fn variance(&self) -> Option<f64> {
Some(self.r * (1.0 - self.p) / (self.p * self.p))
}
/// Returns the skewness of the negative binomial distribution
/// Returns the skewness of the negative binomial distribution.
///
/// # Formula
///
Expand All @@ -199,7 +211,7 @@ impl DiscreteDistribution<f64> for NegativeBinomial {
}

impl Mode<Option<f64>> for NegativeBinomial {
/// Returns the mode for the negative binomial distribution
/// Returns the mode for the negative binomial distribution.
///
/// # Formula
///
Expand All @@ -221,30 +233,50 @@ impl Mode<Option<f64>> for NegativeBinomial {

impl Discrete<u64, f64> for NegativeBinomial {
/// Calculates the probability mass function for the negative binomial
/// distribution at `x`
/// distribution at `x`.
///
/// # Formula
///
/// When `r` is an integer, the formula is:
///
/// ```ignore
/// (x + r - 1 choose x) * (1 - p)^x * p^r
/// ```
///
/// The general formula for real `r` is:
///
/// ```ignore
/// (x + r - 1 choose k) * (1 - p)^x * p^r
/// Γ(r + x)/(Γ(r) * Γ(x + 1)) * (1 - p)^x * p^r
/// ```
///
/// where Γ(x) is the Gamma function.
fn pmf(&self, x: u64) -> f64 {
self.ln_pmf(x).exp()
}

/// Calculates the log probability mass function for the negative binomial
/// distribution at `x`
/// distribution at `x`.
///
/// # Formula
///
/// When `r` is an integer, the formula is:
///
/// ```ignore
/// ln(x + r - 1 choose k) * (1 - p)^x * p^r))
/// ln((x + r - 1 choose x) * (1 - p)^x * p^r)
/// ```
///
/// The general formula for real `r` is:
///
/// ```ignore
/// ln(Γ(r + x)/(Γ(r) * Γ(x + 1)) * (1 - p)^x * p^r)
/// ```
///
/// where Γ(x) is the Gamma function.
fn ln_pmf(&self, x: u64) -> f64 {
let k = x as f64;
gamma::ln_gamma(self.r + k) - gamma::ln_gamma(self.r) - gamma::ln_gamma(k + 1.0)
+ (self.r * self.p.ln())
+ (k * (1.0 - self.p).ln())
+ (k * (-self.p).ln_1p())
}
}

Expand All @@ -254,6 +286,7 @@ mod tests {
use std::fmt::Debug;
use crate::statistics::*;
use crate::distribution::{DiscreteCDF, Discrete, NegativeBinomial};
use crate::distribution::internal::test;
use crate::consts::ACC;

fn try_create(r: f64, p: f64) -> NegativeBinomial {
Expand Down Expand Up @@ -458,19 +491,15 @@ mod tests {
test_case(3.0, 0.5, 1.0, cdf(100));
}

#[test]
fn test_discrete() {
test::check_discrete_distribution(&try_create(5.0, 0.3), 35);
test::check_discrete_distribution(&try_create(10.0, 0.7), 21);
}

#[test]
fn test_sf_upper_bound() {
let sf = |arg: u64| move |x: NegativeBinomial| x.sf(arg);
test_almost(3.0, 0.5, 5.282409836586059e-28, 1e-28, sf(100));
}

// TODO: figure out the best way to re-implement this test. We currently
// do not have a good way to characterize a discrete distribution with a
// CDF that is continuous
//
// #[test]
// fn test_discrete() {
// test::check_discrete_distribution(&try_create(5.0, 0.3), 35);
// test::check_discrete_distribution(&try_create(10.0, 0.7), 21);
// }
}

0 comments on commit 184367c

Please sign in to comment.