diff --git a/src/distribution/negative_binomial.rs b/src/distribution/negative_binomial.rs index fff79190..4da25b1d 100644 --- a/src/distribution/negative_binomial.rs +++ b/src/distribution/negative_binomial.rs @@ -311,10 +311,14 @@ impl Discrete for NegativeBinomial { /// /// 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 * (-self.p).ln_1p()) + if self.p == 1.0 { + if x == 0 { 0.0 } else { f64::NEG_INFINITY } + } else { + 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 * (-self.p).ln_1p()) + } } } @@ -398,17 +402,23 @@ mod tests { test_absolute(1.0, 0.3, 0.3, 1e-15, pmf(0)); test_absolute(1.0, 0.3, 0.21, 1e-15, pmf(1)); test_absolute(3.0, 0.3, 0.027, 1e-15, pmf(0)); + } + + #[test] + fn test_pmf_p_one() { + let pmf = |arg: u64| move |x: NegativeBinomial| x.pmf(arg); + test_exact(1.0, 1.0, 1.0, pmf(0)); + test_exact(1.0, 1.0, 0.0, pmf(1)); + test_exact(2.5, 1.0, 1.0, pmf(0)); + test_exact(2.5, 1.0, 0.0, pmf(1)); + test_exact(0.3, 1.0, 1.0, pmf(0)); test_exact(0.3, 1.0, 0.0, pmf(1)); test_exact(0.3, 1.0, 0.0, pmf(3)); - test_is_nan(0.3, 1.0, pmf(0)); - test_exact(0.3, 1.0, 0.0, pmf(1)); test_exact(0.3, 1.0, 0.0, pmf(10)); - test_is_nan(1.0, 1.0, pmf(0)); - test_exact(1.0, 1.0, 0.0, pmf(1)); - test_is_nan(3.0, 1.0, pmf(0)); + test_exact(3.0, 1.0, 1.0, pmf(0)); test_exact(3.0, 1.0, 0.0, pmf(1)); test_exact(3.0, 1.0, 0.0, pmf(3)); - test_is_nan(10.0, 1.0, pmf(0)); + test_exact(10.0, 1.0, 1.0, pmf(0)); test_exact(10.0, 1.0, 0.0, pmf(1)); test_exact(10.0, 1.0, 0.0, pmf(10)); } @@ -427,17 +437,23 @@ mod tests { test_absolute(1.0, 0.3, -1.203972804, 1e-08, ln_pmf(0)); test_absolute(1.0, 0.3, -1.560647748, 1e-08, ln_pmf(1)); test_absolute(3.0, 0.3, -3.611918413, 1e-08, ln_pmf(0)); + } + + #[test] + fn test_ln_pmf_p_one() { + let ln_pmf = |arg: u64| move |x: NegativeBinomial| x.ln_pmf(arg); + test_exact(1.0, 1.0, 0.0, ln_pmf(0)); + test_exact(1.0, 1.0, f64::NEG_INFINITY, ln_pmf(1)); + test_exact(2.5, 1.0, 0.0, ln_pmf(0)); + test_exact(2.5, 1.0, f64::NEG_INFINITY, ln_pmf(1)); + test_exact(0.3, 1.0, 0.0, ln_pmf(0)); test_exact(0.3, 1.0, f64::NEG_INFINITY, ln_pmf(1)); test_exact(0.3, 1.0, f64::NEG_INFINITY, ln_pmf(3)); - test_is_nan(0.3, 1.0, ln_pmf(0)); - test_exact(0.3, 1.0, f64::NEG_INFINITY, ln_pmf(1)); test_exact(0.3, 1.0, f64::NEG_INFINITY, ln_pmf(10)); - test_is_nan(1.0, 1.0, ln_pmf(0)); - test_exact(1.0, 1.0, f64::NEG_INFINITY, ln_pmf(1)); - test_is_nan(3.0, 1.0, ln_pmf(0)); + test_exact(3.0, 1.0, 0.0, ln_pmf(0)); test_exact(3.0, 1.0, f64::NEG_INFINITY, ln_pmf(1)); test_exact(3.0, 1.0, f64::NEG_INFINITY, ln_pmf(3)); - test_is_nan(10.0, 1.0, ln_pmf(0)); + test_exact(10.0, 1.0, 0.0, ln_pmf(0)); test_exact(10.0, 1.0, f64::NEG_INFINITY, ln_pmf(1)); test_exact(10.0, 1.0, f64::NEG_INFINITY, ln_pmf(10)); }