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
108
109

use crate::{
    rng::{fixup, i2_32m1},
    traits::RNG,
};

const N: usize = 624;
static M: usize = 397;
static MATRIX_A: u32 = 0x9908b0df;
static UPPER_MASK: u32 = 0x80000000;
static LOWER_MASK: u32 = 0x7fffffff;

#[derive(Clone, Debug)]
pub struct MersenneTwister {
    state: [u32; N + 1],
}

impl Default for MersenneTwister {
    fn default() -> Self {
        Self::new()
    }
}

impl MersenneTwister {
    pub fn new() -> Self {
        let mut ret = Self { state: [0; N + 1] };
        ret.randomize();
        ret
    }

    fn mt_sgenrand(&mut self, mut seed: u32, mti: &mut u32) {
        for i in 1..(N + 1) {
            self.state[i] = seed & 0xffff0000;
            seed = seed.wrapping_mul(69069).wrapping_add(1);
            self.state[i] |= (seed & 0xffff0000) >> 16;
            seed = seed.wrapping_mul(69069).wrapping_add(1);
        }
        *mti = N as u32;
    }

    fn mt_genrand(&mut self) -> f64 {
        let mut y: u32 = 0;
        static mag01: [u32; 2] = [0, MATRIX_A];
        /* mag01[x] = x * MATRIX_A  for x=0,1 */
        let mut mti = self.state[0] as u32;
        if mti >= N as u32 {
            /* generate N words at one time */
            if mti == (N + 1) as u32 {
                /* if sgenrand() has not been called, */
                self.mt_sgenrand(4357, &mut mti);
            }

            for kk in 1..((N - M) + 1) {
                y = (self.state[kk] & UPPER_MASK) | (self.state[kk + 1] & LOWER_MASK);
                self.state[kk] = self.state[kk + M] ^ (y >> 1) ^ mag01[(y & 0x1) as usize];
            }
            for kk in ((N - M) + 1)..N {
                y = (self.state[kk] & UPPER_MASK) | (self.state[kk + 1] & LOWER_MASK);

                self.state[kk] = self.state[(kk as i32 + (M as i32 - N as i32)) as usize]
                    ^ (y >> 1)
                    ^ mag01[(y & 0x1) as usize];
            }

            y = (self.state[1] & UPPER_MASK) | (self.state[N] & LOWER_MASK);

            self.state[N] = self.state[M] ^ (y >> 1) ^ mag01[(y & 0x1) as usize];

            mti = 0;
        }
        mti += 1;
        y = self.state[mti as usize];
        y ^= y >> 11;
        y ^= (y << 7) & 0x9d2c5680;
        y ^= (y << 15) & 0xefc60000;
        y ^= y >> 18;
        self.state[0] = mti as u32;
        return y as f64 * i2_32m1;
    }

    pub fn fixup_seeds(&mut self) {
        self.state[0] = N as u32;
        if self.state.iter().all(|&j| j == 0) {
            self.randomize();
        }
    }
}

impl RNG for MersenneTwister {
    fn set_seed(&mut self, mut seed: u32) {
        for _ in 0..50 {
            seed = seed.wrapping_mul(69069).wrapping_add(1);
        }

        (0..(N + 1)).for_each(|j| {
            seed = seed.wrapping_mul(69069).wrapping_add(1);
            self.state[j] = seed;
        });

        self.fixup_seeds()
    }

    fn get_state(&self) -> Vec<i32> {
        self.state.clone().iter().map(|&i| i as i32).collect()
    }

    fn unif_rand(&mut self) -> f64 {
        fixup(self.mt_genrand())
    }
}