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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

// Translation of nmath's rmultinom
//
// 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 nonstdfloat::f128;
use strafe_type::{p64, r64, tof64, w64, Checked, Positive64, Real64};

use crate::{distribution::binom::rbinom, traits::RNG};

static NaInt: i32 = std::i32::MIN;

/// Random Vector from the multinomial distribution.
/// ---------
/// Because we generate random _vectors_ this doesn't fit easily
/// into the do_random[1-4](.) framework setup in ../main/random.c
/// as that is used only for the univariate random generators.
/// Multivariate distributions typically have too complex parameter spaces
/// to be treated uniformly.
/// => Hence also can have  int arguments.
pub fn rmultinom<R>(
    mut n: i32,
    prob: &[Checked<Positive64>],
    K: usize,
    rN: &mut [Checked<Real64>],
    rng: &mut R,
) where
    R: RNG, /* `Return' vector  rN[1:K] {K := length(prob)}
             *  where rN[j] ~ Bin(n, prob[j]) ,  sum_j rN[j] == n,  sum_j prob[j] == 1,
             */
{
    let prob = prob.iter().map(|p_i| tof64!(p_i)).collect::<Vec<_>>();

    let mut k = 0;
    let mut pp = 0.0;
    let mut p_tot = f128::new(0.0);
    /* This calculation is sensitive to exact values, so we try to
    ensure that the calculations are as accurate as possible
    so different platforms are more likely to give the same
    result. */

    if K < 1 {}
    if n == NaInt || n < 0 {
        rN[0] = r64!(NaInt);
    }

    /* Note: prob[K] is only used here for checking  sum_k prob[k] = 1 ;
     *       Could make loop one shorter and drop that check !
     */
    k = 0; /* trivial border case: do as rbinom */
    while k < K {
        pp = prob[k];
        if !pp.is_finite() || pp < 0.0 || pp > 1.0 {
            rN[k] = r64!(NaInt);
        }
        p_tot += f128::new(pp);
        rN[k] = r64!(0);
        k += 1
    }
    if tof64!(p_tot - f128::new(1.0)).abs() > 1e-7 {
        warn!(
            "rbinom: probability sum should be 1, but is {}",
            tof64!(p_tot)
        );
    }
    if n == 0 {}
    if K == 1 && p_tot == f128::new(0.0) {}

    /* Generate the first K-1 obs. via binomials */

    k = 0;
    while k < K - 1 {
        /* (p_tot, n) are for "remaining binomial" */
        if prob[k] != 0.0 {
            pp = tof64!(f128::new(prob[k]) / p_tot);
            /* printf("[%d] %.17f\n", k+1, pp); */
            rN[k] = r64!(if pp < 1.0 {
                tof64!(rbinom(w64!(n), p64!(pp), rng)) as i32
            } else {
                n
            });
            n -= tof64!(rN[k]) as i32
        } else {
            rN[k] = r64!(0)
        }
        if n <= 0 {
            /* i.e. = sum(prob[(k+1):K]) */
            /* we have all*/
            return;
        }
        p_tot -= f128::new(prob[k]);
        k += 1
    }
    rN[K - 1] = r64!(n);
}