vosealias/
lib.rs

1//! # Walker-Vose Alias Method
2//! A simple implementation of alias tables using the Walker-Vose method.
3
4extern crate num_traits;
5extern crate rand;
6
7use std::fmt;
8use std::iter::{FromIterator, Sum};
9use std::vec::Vec;
10
11use num_traits::{Float, NumCast, One, Zero};
12
13use rand::Rng;
14use rand::distributions::range::{Range, SampleRange};
15use rand::distributions::IndependentSample;
16
17
18#[derive(Debug)]
19enum AliasEntry<F> {
20    Aliased {
21        threshold: F,
22        value: usize,
23        alias: usize,
24    },
25    Unaliased(usize),
26}
27use AliasEntry::*;
28
29
30/// An alias table, which uses floating point probabilities of type `F` and table entries of type
31/// `T`.
32pub struct AliasTable<T, F> {
33    table: Vec<AliasEntry<F>>,
34    objs: Vec<T>,
35    range: Range<usize>,
36    float: Range<F>,
37}
38
39/// An iterator for an alias table.
40#[derive(Clone)]
41pub struct AliasTableIterator<'a, T: 'a, F: 'a, R>
42    where R: Rng + Sized
43{
44    rng: R,
45    table: &'a AliasTable<T, F>
46}
47
48
49impl<T, F> fmt::Debug for AliasTable<T, F>
50    where F: fmt::Debug
51{
52    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
53        write!(fmt, "AliasTable {{ table: {:?} }}", self.table)
54    }
55}
56
57
58impl<T, F> AliasTable<T, F>
59    where F: PartialOrd + SampleRange
60{
61    /// Pick a random element from the distribution. Samples from the RNG using `ind_sample` only.
62    pub fn pick<'a, R: Rng>(&'a self, rng: &mut R) -> &'a T {
63        let idx = self.range.ind_sample(rng);
64        let entry = &self.table[idx];
65        match *entry {
66            Aliased { ref threshold, value, alias } => {
67                if &self.float.ind_sample(rng) < threshold {
68                    &self.objs[value]
69                } else {
70                    &self.objs[alias]
71                }
72            }
73            Unaliased(idx) => &self.objs[idx],
74        }
75    }
76
77    /// Given an RNG, produce an iterator that picks random element from the distribution by
78    /// calling `pick` repeatedly with the given RNG.
79    pub fn iter<R: Rng>(&self, rng: R) -> AliasTableIterator<T, F, R> {
80        AliasTableIterator {
81            rng: rng,
82            table: self
83        }
84    }
85}
86
87impl<'a, T, F: 'a> FromIterator<(T, F)> for AliasTable<T, F>
88    where F: Float + NumCast + One + SampleRange + Sum<F> + Zero
89{
90    /// Construct an alias table from an iterator. Expects a tuple, where the left-hand element is
91    /// the distribution's value, and the right-hand element is the value's weight in the distribution.
92    fn from_iter<I: IntoIterator<Item = (T, F)>>(iter: I) -> Self {
93        let (objs, ps): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
94        let psum: F = ps.iter().cloned().sum();
95
96        let pn = F::from(ps.len())
97            .expect("Error casting usize to generic parameter F of AliasTable<T, F>");
98        let pcoeff = pn / psum;
99
100        let (mut small, mut large): (Vec<_>, Vec<_>) =
101            ps.into_iter().map(|p| pcoeff * p).enumerate().partition(|&(_, p)| p < F::one());
102        let mut table = Vec::new();
103
104
105        while !(small.is_empty() || large.is_empty()) {
106            let (l, p_l) = small.pop().unwrap();
107            let (g, p_g) = large.pop().unwrap();
108
109            table.push(Aliased {
110                threshold: p_l,
111                value: l,
112                alias: g,
113            });
114
115            let p_g = (p_g + p_l) - F::one();
116
117            if p_g < F::one() {
118                    &mut small
119                } else {
120                    &mut large
121                }
122                .push((g, p_g));
123        }
124
125        table.extend(large.iter().map(|&(g, _)| Unaliased(g)));
126
127        table.extend(small.iter().map(|&(l, _)| Unaliased(l)));
128
129        AliasTable {
130            range: Range::new(0, table.len()),
131            float: Range::new(F::zero(), F::one()),
132            table: table,
133            objs: objs,
134        }
135    }
136}
137
138impl<'a, T: 'a, F, R> Iterator for AliasTableIterator<'a, T, F, R>
139    where F: PartialOrd + SampleRange,
140          R: Rng
141{
142    type Item = &'a T;
143
144    fn next(&mut self) -> Option<Self::Item> {
145        Some(self.table.pick(&mut self.rng))
146    }
147
148    fn size_hint(&self) -> (usize, Option<usize>) {
149        (std::usize::MAX, None)
150    }
151}
152
153impl<'a, T, F> IntoIterator for &'a AliasTable<T, F>
154    where F: Sized + PartialOrd + SampleRange
155{
156    type Item = &'a T;
157    type IntoIter = AliasTableIterator<'a, T, F, rand::ThreadRng>;
158
159    /// Produces an iterator that picks random element from the distribution by calling
160    /// calling `pick` repeatedly with the thread's RNG.
161    fn into_iter(self) -> Self::IntoIter {
162        self.iter(rand::thread_rng())
163    }
164}
vosealias/lib.rs

vosealias/
lib.rs
