bcur 2.0.0

A Rust implementation of Uniform Resources
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
110
111

//! Walker's alias method sampler (ordered partition matching `ur-rs`).
//!
//! Indexing uses known-valid table indices produced by the algorithm; panicking
//! on OOB would indicate a logic bug rather than hostile input.

#![allow(
    clippy::indexing_slicing,
    reason = "alias-table construction uses closed-form indices 0..count matching ur-rs"
)]

use alloc::vec::Vec;

use super::Xoshiro256;

/// Weighted discrete sampler over `0..weights.len()`.
#[derive(Debug)]
pub(crate) struct Sampler {
    aliases: Vec<u32>,
    probs: Vec<f64>,
}

#[allow(
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss,
    reason = "alias method uses f64 probabilities; table sizes match ur-rs exactly"
)]
impl Sampler {
    /// Builds a sampler from non-negative weights that sum to a positive value.
    ///
    /// # Panics
    ///
    /// Panics if any weight is negative or the sum is not positive. Callers must
    /// only construct this from valid degree weights (`1/k`).
    pub(crate) fn new(mut weights: Vec<f64>) -> Self {
        assert!(
            !weights.iter().any(|&p| p < 0.0),
            "negative probability encountered"
        );
        let summed = weights.iter().sum::<f64>();
        assert!(summed > 0.0, "probabilities don't sum to a positive value");
        let count = weights.len();
        for w in &mut weights {
            *w *= count as f64 / summed;
        }
        // Ordered partition: indices count-1 down to 0, small then large.
        let (mut s, mut l): (Vec<usize>, Vec<usize>) = (1..=count)
            .map(|j| count - j)
            .partition(|&j| weights[j] < 1.0);

        let mut probs: Vec<f64> = alloc::vec![0.0; count];
        let mut aliases: Vec<u32> = alloc::vec![0; count];

        while !s.is_empty() && !l.is_empty() {
            let a = s.remove(s.len() - 1);
            let g = l.remove(l.len() - 1);
            probs[a] = weights[a];
            aliases[a] = g as u32;
            weights[g] += weights[a] - 1.0;
            if weights[g] < 1.0 {
                s.push(g);
            } else {
                l.push(g);
            }
        }

        while !l.is_empty() {
            let g = l.remove(l.len() - 1);
            probs[g] = 1.0;
        }

        while !s.is_empty() {
            let a = s.remove(s.len() - 1);
            probs[a] = 1.0;
        }

        Self { aliases, probs }
    }

    pub(crate) fn next(&self, xoshiro: &mut Xoshiro256) -> u32 {
        let n = self.probs.len();
        // `next_int` covers the `r1 == 1.0` edge by clamping to `n - 1` (same
        // documented deviation), so `i` is always a valid table index.
        let i = xoshiro.next_index(n);
        let r2 = xoshiro.next_double();
        if r2 < self.probs[i] {
            i as u32
        } else {
            self.aliases[i]
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sampler() {
        let weights = vec![1.0, 2.0, 4.0, 8.0];
        let mut xoshiro = Xoshiro256::from("Wolf");
        let sampler = Sampler::new(weights);
        let expected = [
            3, 3, 3, 3, 3, 3, 3, 0, 2, 3, 3, 3, 3, 1, 2, 2, 1, 3, 3, 2, 3, 3, 1, 1, 2, 1, 1, 3, 1,
            3,
        ];
        for e in expected {
            assert_eq!(sampler.next(&mut xoshiro), e);
        }
    }
}