zsfm_tabfm/ensemble/
pyrandom.rs

1//! A bit-compatible port of CPython's `random.Random` (Mersenne Twister, MT19937), so ensemble
2//! member generation can be validated numerically against the real `TabFMClassifier`/
3//! `TabFMRegressor` sklearn wrapper (which seeds `random.Random(random_state)` and calls
4//! `.sample()`/`.shuffle()` in a specific order — see `config_gen.rs`).
5//!
6//! Ported from CPython's `Modules/_randommodule.c` (MT19937 core, `init_by_array` seeding) and
7//! `Lib/random.py` (`_randbelow`, `sample`, `shuffle`). Reference sequences used in the unit
8//! tests below were generated with the project's own `.venv/bin/python3`.
9
10const N: usize = 624;
11const M: usize = 397;
12const MATRIX_A: u32 = 0x9908_b0df;
13const UPPER_MASK: u32 = 0x8000_0000;
14const LOWER_MASK: u32 = 0x7fff_ffff;
15
16pub struct Mt19937 {
17    mt: [u32; N],
18    index: usize,
19}
20
21impl Mt19937 {
22    fn init_genrand(seed: u32) -> Self {
23        let mut mt = [0u32; N];
24        mt[0] = seed;
25        for i in 1..N {
26            mt[i] = (1_812_433_253u32.wrapping_mul(mt[i - 1] ^ (mt[i - 1] >> 30)))
27                .wrapping_add(i as u32);
28        }
29        Mt19937 { mt, index: N }
30    }
31
32    /// CPython seeds every integer via `init_by_array`, never bare `init_genrand` — the seed
33    /// integer is first split into little-endian 32-bit words (`key`).
34    pub fn from_seed_key(key: &[u32]) -> Self {
35        let mut rng = Self::init_genrand(19_650_218);
36        let key = if key.is_empty() {
37            vec![0u32]
38        } else {
39            key.to_vec()
40        };
41        let key_length = key.len();
42        let mut i = 1usize;
43        let mut j = 0usize;
44        for _ in 0..N.max(key_length) {
45            let prev = rng.mt[i - 1];
46            rng.mt[i] = (rng.mt[i] ^ ((prev ^ (prev >> 30)).wrapping_mul(1_664_525)))
47                .wrapping_add(key[j])
48                .wrapping_add(j as u32);
49            i += 1;
50            j += 1;
51            if i >= N {
52                rng.mt[0] = rng.mt[N - 1];
53                i = 1;
54            }
55            if j >= key_length {
56                j = 0;
57            }
58        }
59        for _ in 0..N - 1 {
60            let prev = rng.mt[i - 1];
61            rng.mt[i] = (rng.mt[i] ^ ((prev ^ (prev >> 30)).wrapping_mul(1_566_083_941)))
62                .wrapping_sub(i as u32);
63            i += 1;
64            if i >= N {
65                rng.mt[0] = rng.mt[N - 1];
66                i = 1;
67            }
68        }
69        rng.mt[0] = 0x8000_0000;
70        rng
71    }
72
73    /// Seed from a small non-negative integer (as CPython does for `random.Random(42)` etc).
74    pub fn from_u64_seed(seed: u64) -> Self {
75        if seed == 0 {
76            return Self::from_seed_key(&[]);
77        }
78        let mut key = Vec::new();
79        let mut s = seed;
80        while s > 0 {
81            key.push((s & 0xffff_ffff) as u32);
82            s >>= 32;
83        }
84        Self::from_seed_key(&key)
85    }
86
87    fn regenerate(&mut self) {
88        let mag01 = [0u32, MATRIX_A];
89        for kk in 0..N - M {
90            let y = (self.mt[kk] & UPPER_MASK) | (self.mt[kk + 1] & LOWER_MASK);
91            self.mt[kk] = self.mt[kk + M] ^ (y >> 1) ^ mag01[(y & 1) as usize];
92        }
93        for kk in N - M..N - 1 {
94            let y = (self.mt[kk] & UPPER_MASK) | (self.mt[kk + 1] & LOWER_MASK);
95            self.mt[kk] = self.mt[kk + M - N] ^ (y >> 1) ^ mag01[(y & 1) as usize];
96        }
97        let y = (self.mt[N - 1] & UPPER_MASK) | (self.mt[0] & LOWER_MASK);
98        self.mt[N - 1] = self.mt[M - 1] ^ (y >> 1) ^ mag01[(y & 1) as usize];
99        self.index = 0;
100    }
101
102    pub fn next_u32(&mut self) -> u32 {
103        if self.index >= N {
104            self.regenerate();
105        }
106        let mut y = self.mt[self.index];
107        self.index += 1;
108        y ^= y >> 11;
109        y ^= (y << 7) & 0x9d2c_5680;
110        y ^= (y << 15) & 0xefc6_0000;
111        y ^= y >> 18;
112        y
113    }
114
115    /// CPython's `getrandbits(k)`: words filled least-significant-first, 32 bits at a time.
116    pub fn getrandbits(&mut self, k: u32) -> u64 {
117        if k <= 32 {
118            return (self.next_u32() >> (32 - k)) as u64;
119        }
120        let words = (k - 1) / 32 + 1;
121        let mut result: u64 = 0;
122        for w in 0..words {
123            let mut r = self.next_u32();
124            if w == words - 1 {
125                r >>= 32 * words - k;
126            }
127            result |= (r as u64) << (32 * w);
128        }
129        result
130    }
131}
132
133/// A CPython-compatible `random.Random` instance: MT19937 state plus the pure-Python wrapper
134/// logic (`_randbelow`, `sample`, `shuffle`) from `Lib/random.py`.
135pub struct PyRandom {
136    mt: Mt19937,
137}
138
139impl PyRandom {
140    pub fn new(seed: u64) -> Self {
141        PyRandom {
142            mt: Mt19937::from_u64_seed(seed),
143        }
144    }
145
146    fn bit_length(n: u64) -> u32 {
147        if n == 0 {
148            0
149        } else {
150            64 - n.leading_zeros()
151        }
152    }
153
154    /// `Random._randbelow_with_getrandbits`: rejection-sampled `getrandbits`.
155    pub fn randbelow(&mut self, n: u64) -> u64 {
156        if n == 0 {
157            return 0;
158        }
159        let k = Self::bit_length(n);
160        loop {
161            let r = self.mt.getrandbits(k);
162            if r < n {
163                return r;
164            }
165        }
166    }
167
168    /// `Random.sample(range(n), k)` restricted to integer populations `0..n` (the only case
169    /// needed here — feature/class/row indices), returning the sampled indices in draw order.
170    pub fn sample_indices(&mut self, n: usize, k: usize) -> Vec<usize> {
171        assert!(k <= n, "sample larger than population");
172        let mut setsize: f64 = 21.0;
173        if k > 5 {
174            setsize += 4f64.powf((3.0 * k as f64).log(4.0).ceil());
175        }
176        let mut result = vec![0usize; k];
177        if (n as f64) <= setsize {
178            let mut pool: Vec<usize> = (0..n).collect();
179            for i in 0..k {
180                let j = self.randbelow((n - i) as u64) as usize;
181                result[i] = pool[j];
182                pool[j] = pool[n - i - 1];
183            }
184        } else {
185            let mut selected = std::collections::HashSet::new();
186            for i in 0..k {
187                let mut j = self.randbelow(n as u64) as usize;
188                while selected.contains(&j) {
189                    j = self.randbelow(n as u64) as usize;
190                }
191                selected.insert(j);
192                result[i] = j;
193            }
194        }
195        result
196    }
197
198    /// `Random.sample(population, k)` for an arbitrary slice, via `sample_indices`.
199    pub fn sample<T: Clone>(&mut self, population: &[T], k: usize) -> Vec<T> {
200        self.sample_indices(population.len(), k)
201            .into_iter()
202            .map(|i| population[i].clone())
203            .collect()
204    }
205
206    /// `Random.shuffle(x)`: in-place Fisher-Yates via `_randbelow`.
207    pub fn shuffle<T>(&mut self, x: &mut [T]) {
208        let len = x.len();
209        if len < 2 {
210            return;
211        }
212        for i in (1..len).rev() {
213            let j = self.randbelow((i + 1) as u64) as usize;
214            x.swap(i, j);
215        }
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn test_sample_10_of_10() {
225        // python3 -c "import random; print(random.Random(42).sample(range(10), 10))"
226        let mut r = PyRandom::new(42);
227        assert_eq!(r.sample_indices(10, 10), vec![1, 0, 4, 9, 6, 5, 8, 2, 3, 7]);
228    }
229
230    #[test]
231    fn test_sample_3_of_5() {
232        // python3 -c "import random; print(random.Random(42).sample(range(5), 3))"
233        let mut r = PyRandom::new(42);
234        assert_eq!(r.sample_indices(5, 3), vec![0, 4, 2]);
235    }
236
237    #[test]
238    fn test_shuffle_8() {
239        // python3 -c "import random; l=list(range(8)); random.Random(42).shuffle(l); print(l)"
240        let mut r = PyRandom::new(42);
241        let mut v: Vec<usize> = (0..8).collect();
242        r.shuffle(&mut v);
243        assert_eq!(v, vec![3, 4, 6, 7, 2, 5, 0, 1]);
244    }
245
246    #[test]
247    fn test_sample_5_of_100_rejection_branch() {
248        // python3 -c "import random; print(random.Random(7).sample(range(100), 5))"
249        // k=5 (not >5) so setsize=21; n=100 > 21, exercises the rejection-set branch.
250        let mut r = PyRandom::new(7);
251        assert_eq!(r.sample_indices(100, 5), vec![41, 19, 50, 83, 6]);
252    }
253}
zsfm_tabfm/ensemble/pyrandom.rs

zsfm_tabfm/ensemble/
pyrandom.rs
