r2rs-nmath 0.1.1

// Translation of nmath's dnbinom
//
// Copyright (C) 2020 neek-sss
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use num_traits::Float;
use strafe_type::{FloatConstraint, Positive64, Probability64, Real64};

use crate::{
    distribution::{binom::dbinom_raw, gamma::lgamma1p, pois::dpois_raw},
    traits::DPQ,
};

/// Computes the negative binomial distribution. For integer n,
/// this is probability of x failures before the nth success in a
/// sequence of Bernoulli trials. We do not enforce integer n, since
/// the distribution is well defined for non-integers,
/// and this can be useful for e.g. overdispersed discrete survival times.
pub fn dnbinom<R: Into<Real64>, PO: Into<Positive64>, PR: Into<Probability64>>(
    x: R,
    size: PO,
    prob: PR,
    log: bool,
) -> Real64 {
    let mut x = x.into().unwrap();
    let mut size = size.into().unwrap();
    let prob = prob.into().unwrap();

    let mut ans = 0.0;

    if prob <= 0.0 {
        return f64::nan().into();
    }
    if x.is_non_integer() {
        warn!("non-integer x = {}", x);
        return f64::d_0(log).into();
    }
    if x < 0.0 || !x.is_finite() {
        return f64::d_0(log).into();
    }
    if x == 0.0 {
        /* limiting case as size approaches zero is point mass at zero */
        return if size == 0.0 {
            f64::d_1(log).into()
        } else if log {
            // size > 0:  P(x, ..) = pr^n :
            (size * prob.ln()).into()
        } else {
            prob.powf(size).into()
        };
    }
    x = x.round();
    if !size.is_finite() {
        size = f64::MAX
    }

    if x < 1e-10 * size {
        // instead of dbinom_raw(), use 2 terms of Abramowitz & Stegun (6.1.47)
        (size * prob.ln() + x * (size.ln() + (-prob).ln_1p()) - lgamma1p(x)
            + (x * (x - 1.0) / (2.0 * size)).ln_1p())
        .d_exp(log)
        .into()
    } else {
        /* log( size/(size+x) ) is much less accurate than log1p(- x/(size+x))
        for |x| << size (and actually when x < size): */
        let p = if log {
            if x < size {
                (-x / (size + x)).ln_1p()
            } else {
                (size / (size + x)).ln()
            }
        } else {
            size / (size + x)
        };
        ans = dbinom_raw(size, x + size, prob, 1.0 - prob, log);
        if log { p + ans } else { p * ans }.into()
    }
}

pub fn dnbinom_mu<R: Into<Real64>, PO1: Into<Positive64>, PO2: Into<Positive64>>(
    x: R,
    size: PO1,
    mu: PO2,
    log: bool,
) -> Real64 {
    let mut x = x.into().unwrap();
    let size = size.into().unwrap();
    let mu = mu.into().unwrap();

    /* originally, just set  prob :=  size / (size + mu)  and called dbinom_raw(),
     * but that suffers from cancellation when   mu << size  */

    if x.is_non_integer() {
        warn!("non-integer x = {}", x);
        return f64::d_0(log).into();
    }
    if x < 0.0 || !x.is_finite() {
        return f64::d_0(log).into();
    }

    /* limiting case as size approaches zero is point mass at zero,
     * even if mu is kept constant. limit distribution does not
     * have mean mu, though.
     */
    if x == 0.0 && size == 0.0 {
        return f64::d_1(log).into();
    }
    x = x.round();
    // FIXME use also for size "almost" Inf because that gives NaN ???
    if !size.is_finite() {
        // limit case: Poisson
        return dpois_raw(x, mu, log).into();
    }

    if x == 0.0 {
        /* be accurate, both for n << mu, and n >> mu :*/
        (size
            * if size < mu {
                (size / (size + mu)).ln()
            } else {
                (-mu / (size + mu)).ln_1p()
            })
        .d_exp(log);
    }
    if x < 1e-10 * size {
        /* don't use dbinom_raw() but MM's formula: */
        /* FIXME --- 1e-8 shows problem; rather use algdiv() from ./toms708.c */
        let p = if size < mu {
            (size / (1.0 + size / mu)).ln()
        } else {
            (mu / (1.0 + mu / size)).ln()
        };
        ((x * p - mu - lgamma1p(x)) + (x * (x - 1.0) / (2.0 * size)).ln_1p()).d_exp(log)
    } else {
        /* no unnecessary cancellation inside dbinom_raw, when
        x_ = size and n_ = x+size are so close that n_ - x_ loses accuracy
        but log( size/(size+x) ) is much less accurate than log1p(- x/(size+x))
        for |x| << size (and actually when x < size): */
        let p_0 = if log {
            if x < size {
                (-x / (size + x)).ln_1p()
            } else {
                (size / (size + x)).ln()
            }
        } else {
            size / (size + x)
        };
        let ans = dbinom_raw(size, x + size, size / (size + mu), mu / (size + mu), log);
        if log {
            p_0 + ans
        } else {
            p_0 * ans
        }
    }
    .into()
}