r2rs-nmath 0.1.1

Statistics programming for Rust based on R's nmath package
Documentation
// Translation of nmath's phyper
//
// 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 std::mem::swap;

use nonstdfloat::f128;
use num_traits::{Float, ToPrimitive};
use strafe_type::{FloatConstraint, LogProbability64, PositiveInteger64, Probability64, Real64};

use crate::{distribution::hyper::dhyper, traits::DPQ};

/// Calculate
///
/// $ log(\frac{phyper (x, NR, NB, n, true, false)}{dhyper (x, NR, NB, n, false)})
///
/// without actually calling phyper.  This assumes that
///
/// $ x * (NR + NB) <= n * NR $
fn pdhyper(mut x: f64, NR: f64, NB: f64, n: f64, log: bool) -> f64 {
    let mut sum = f128::new(0.0);
    let mut term = f128::new(1.0);

    while x > 0.0 && term >= f128::new(2.2204460492503131e-16) * sum {
        term *= f128::new(x * (NB - n + x) / (n + 1.0 - x) / (NR + 1.0 - x));
        sum += term;
        x -= 1.
    }

    let ss = sum.to_f64().unwrap();
    // Fix some NAN problems
    if log {
        if ss < -1.0 { -1.0 } else { ss }.ln_1p()
    } else {
        1.0 + ss
    }
}

/// The distribution function of the hypergeometric distribution.
/// Sample of  n balls from  NR red  and  NB black ones;  x are red
pub fn log_phyper<
    R: Into<Real64>,
    P1: Into<PositiveInteger64>,
    P2: Into<PositiveInteger64>,
    P3: Into<PositiveInteger64>,
>(
    x: R,
    NR: P1,
    NB: P2,
    n: P3,
    lower_tail: bool,
) -> LogProbability64 {
    phyper_inner(x, NR, NB, n, lower_tail, true).into()
}

/// The distribution function of the hypergeometric distribution.
/// Sample of  n balls from  NR red  and  NB black ones;  x are red
pub fn phyper<
    R: Into<Real64>,
    P1: Into<PositiveInteger64>,
    P2: Into<PositiveInteger64>,
    P3: Into<PositiveInteger64>,
>(
    x: R,
    NR: P1,
    NB: P2,
    n: P3,
    lower_tail: bool,
) -> Probability64 {
    phyper_inner(x, NR, NB, n, lower_tail, false).into()
}

pub fn phyper_inner<
    R: Into<Real64>,
    P1: Into<PositiveInteger64>,
    P2: Into<PositiveInteger64>,
    P3: Into<PositiveInteger64>,
>(
    x: R,
    NR: P1,
    NB: P2,
    n: P3,
    mut lower_tail: bool,
    log: bool,
) -> f64 {
    let mut x = x.into().unwrap();
    let mut NR = NR.into().unwrap();
    let mut NB = NB.into().unwrap();
    let mut n = n.into().unwrap();

    let mut d = 0.0;
    let mut pd = 0.0;

    x = (x + 1e-7).floor();
    NR = NR.round();
    NB = NB.round();
    n = n.round();

    if !NR.is_finite() || !NB.is_finite() || n > NR + NB {
        return f64::nan();
    }

    if x * (NR + NB) > n * NR {
        /* Swap tails. */
        swap(&mut NB, &mut NR);
        x = n - x - 1.0;
        lower_tail = !lower_tail
    }

    /* support of dhyper() as a function of its parameters
     * R:  .suppHyper <- function(m,n,k) max(0, k-n) : min(k, m)
     * --  where R's (m,n, k) == (NR,NB, n)  here */
    if x < 0.0 || x < n - NB {
        return f64::dt_0(lower_tail, log);
    }
    if x >= NR || x >= n {
        return f64::dt_1(lower_tail, log);
    }

    d = dhyper(x, NR, NB, n, log).unwrap();

    // dhyper(.., log_p=FALSE) > 0 mathematically, but not always numerically :
    if (!log && d == 0.0) || (log && d == f64::neg_infinity()) {
        return f64::dt_0(lower_tail, log);
    }

    pd = pdhyper(x, NR, NB, n, log);

    if log {
        (d + pd).dt_log(lower_tail, true)
    } else {
        (d * pd).d_lval(lower_tail)
    }
}

// NB: MM has code for  AS 152 (Lund, 1980) >> R_77 (Shea, 1989) >> R_86 (Berger, 1991)