r2rs-nmath 0.1.1

Statistics programming for Rust based on R's nmath package
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

// Translation of nmath's dhyper
//
// 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, PositiveInteger64, Real64};

use crate::{distribution::binom::dbinom_raw, traits::DPQ};

/// Given a sequence of r successes and b failures, we sample n (\le b+r)
/// items without replacement. The hypergeometric probability is the
/// probability of x successes:
///
/// $ p(x; r,b,n) = \frac{choose(r, x) * choose(b, n-x)}{choose(r+b, n)} = \frac{dbinom(x,r,p) * dbinom(n-x,b,p)}{dbinom(n,r+b,p)} $
///
/// for any p. For numerical stability, we take $p=\frac{n}{r+b}$; with this choice,
/// the denominator is not exponentially small.
pub fn dhyper<
    R: Into<Real64>,
    P1: Into<PositiveInteger64>,
    P2: Into<PositiveInteger64>,
    P3: Into<PositiveInteger64>,
>(
    x: R,
    r: P1,
    b: P2,
    n: P3,
    log: bool,
) -> Real64 {
    let mut x = x.into().unwrap();
    let mut r = r.into().unwrap();
    let mut b = b.into().unwrap();
    let mut n = n.into().unwrap();

    let mut p = 0.0;
    let mut q = 0.0;
    let mut p1 = 0.0;
    let mut p2 = 0.0;
    let mut p3 = 0.0;

    if n > r + b {
        return f64::nan().into();
    }

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

    x = x.round();
    r = r.round();
    b = b.round();
    n = n.round();

    if n < x || r < x || n - x > b {
        return f64::d_0(log).into();
    }
    if n == 0.0 {
        return if x == 0.0 {
            f64::d_1(log)
        } else {
            f64::d_0(log)
        }
        .into();
    }

    p = n / (r + b);
    q = (r + b - n) / (r + b);

    p1 = dbinom_raw(x, r, p, q, log);
    p2 = dbinom_raw(n - x, b, p, q, log);
    p3 = dbinom_raw(n, r + b, p, q, log);

    if log { (p1 + p2) - p3 } else { (p1 * p2) / p3 }.into()
}