1 // Copyright 2013 The Rust Project Developers. See the COPYRIGHT
2 // file at the top-level directory of this distribution and at
3 // http://rust-lang.org/COPYRIGHT.
4 //
5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8 // option. This file may not be copied, modified, or distributed
9 // except according to those terms.
10
11 /*!
12 Sampling from random distributions.
13
14 This is a generalization of `Rand` to allow parameters to control the
15 exact properties of the generated values, e.g. the mean and standard
16 deviation of a normal distribution. The `Sample` trait is the most
17 general, and allows for generating values that change some state
18 internally. The `IndependentSample` trait is for generating values
19 that do not need to record state.
20
21 */
22
23 use std::num;
24 use std::num::CheckedAdd;
25 use {Rng, Rand};
26
27 pub use self::range::Range;
28 pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
29 pub use self::normal::{Normal, LogNormal};
30 pub use self::exponential::Exp;
31
32 pub mod range;
33 pub mod gamma;
34 pub mod normal;
35 pub mod exponential;
36
37 /// Types that can be used to create a random instance of `Support`.
38 pub trait Sample<Support> {
39 /// Generate a random value of `Support`, using `rng` as the
40 /// source of randomness.
41 fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
42 }
43
44 /// `Sample`s that do not require keeping track of state.
45 ///
46 /// Since no state is recorded, each sample is (statistically)
47 /// independent of all others, assuming the `Rng` used has this
48 /// property.
49 // FIXME maybe having this separate is overkill (the only reason is to
50 // take &self rather than &mut self)? or maybe this should be the
51 // trait called `Sample` and the other should be `DependentSample`.
52 pub trait IndependentSample<Support>: Sample<Support> {
53 /// Generate a random value.
54 fn ind_sample<R: Rng>(&self, &mut R) -> Support;
55 }
56
57 /// A wrapper for generating types that implement `Rand` via the
58 /// `Sample` & `IndependentSample` traits.
59 pub struct RandSample<Sup>;
60
61 impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
62 fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
63 }
64
65 impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
66 fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
67 rng.gen()
68 }
69 }
70
71 /// A value with a particular weight for use with `WeightedChoice`.
72 pub struct Weighted<T> {
73 /// The numerical weight of this item
74 pub weight: uint,
75 /// The actual item which is being weighted
76 pub item: T,
77 }
78
79 /// A distribution that selects from a finite collection of weighted items.
80 ///
81 /// Each item has an associated weight that influences how likely it
82 /// is to be chosen: higher weight is more likely.
83 ///
84 /// The `Clone` restriction is a limitation of the `Sample` and
85 /// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
86 /// all `T`, as is `uint`, so one can store references or indices into
87 /// another vector.
88 ///
89 /// # Example
90 ///
91 /// ```rust
92 /// use rand::distributions::{Weighted, WeightedChoice, IndependentSample};
93 ///
94 /// let wc = WeightedChoice::new(vec!(Weighted { weight: 2, item: 'a' },
95 /// Weighted { weight: 4, item: 'b' },
96 /// Weighted { weight: 1, item: 'c' }));
97 /// let mut rng = rand::task_rng();
98 /// for _ in range(0, 16) {
99 /// // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
100 /// println!("{}", wc.ind_sample(&mut rng));
101 /// }
102 /// ```
103 pub struct WeightedChoice<T> {
104 items: Vec<Weighted<T>>,
105 weight_range: Range<uint>
106 }
107
108 impl<T: Clone> WeightedChoice<T> {
109 /// Create a new `WeightedChoice`.
110 ///
111 /// Fails if:
112 /// - `v` is empty
113 /// - the total weight is 0
114 /// - the total weight is larger than a `uint` can contain.
115 pub fn new(mut items: Vec<Weighted<T>>) -> WeightedChoice<T> {
116 // strictly speaking, this is subsumed by the total weight == 0 case
117 assert!(!items.is_empty(), "WeightedChoice::new called with no items");
118
119 let mut running_total = 0u;
120
121 // we convert the list from individual weights to cumulative
122 // weights so we can binary search. This *could* drop elements
123 // with weight == 0 as an optimisation.
124 for item in items.mut_iter() {
125 running_total = running_total.checked_add(&item.weight)
126 .expect("WeightedChoice::new called with a total weight larger \
127 than a uint can contain");
128
129 item.weight = running_total;
130 }
131 assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");
132
133 WeightedChoice {
134 items: items,
135 // we're likely to be generating numbers in this range
136 // relatively often, so might as well cache it
137 weight_range: Range::new(0, running_total)
138 }
139 }
140 }
141
142 impl<T: Clone> Sample<T> for WeightedChoice<T> {
143 fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
144 }
145
146 impl<T: Clone> IndependentSample<T> for WeightedChoice<T> {
147 fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
148 // we want to find the first element that has cumulative
149 // weight > sample_weight, which we do by binary since the
150 // cumulative weights of self.items are sorted.
151
152 // choose a weight in [0, total_weight)
153 let sample_weight = self.weight_range.ind_sample(rng);
154
155 // short circuit when it's the first item
156 if sample_weight < self.items.get(0).weight {
157 return self.items.get(0).item.clone();
158 }
159
160 let mut idx = 0;
161 let mut modifier = self.items.len();
162
163 // now we know that every possibility has an element to the
164 // left, so we can just search for the last element that has
165 // cumulative weight <= sample_weight, then the next one will
166 // be "it". (Note that this greatest element will never be the
167 // last element of the vector, since sample_weight is chosen
168 // in [0, total_weight) and the cumulative weight of the last
169 // one is exactly the total weight.)
170 while modifier > 1 {
171 let i = idx + modifier / 2;
172 if self.items.get(i).weight <= sample_weight {
173 // we're small, so look to the right, but allow this
174 // exact element still.
175 idx = i;
176 // we need the `/ 2` to round up otherwise we'll drop
177 // the trailing elements when `modifier` is odd.
178 modifier += 1;
179 } else {
180 // otherwise we're too big, so go left. (i.e. do
181 // nothing)
182 }
183 modifier /= 2;
184 }
185 return self.items.get(idx + 1).item.clone();
186 }
187 }
188
189 mod ziggurat_tables;
190
191 /// Sample a random number using the Ziggurat method (specifically the
192 /// ZIGNOR variant from Doornik 2005). Most of the arguments are
193 /// directly from the paper:
194 ///
195 /// * `rng`: source of randomness
196 /// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
197 /// * `X`: the $x_i$ abscissae.
198 /// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
199 /// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
200 /// * `pdf`: the probability density function
201 /// * `zero_case`: manual sampling from the tail when we chose the
202 /// bottom box (i.e. i == 0)
203
204 // the perf improvement (25-50%) is definitely worth the extra code
205 // size from force-inlining.
206 #[inline(always)]
207 fn ziggurat<R:Rng>(
208 rng: &mut R,
209 symmetric: bool,
210 x_tab: ziggurat_tables::ZigTable,
211 f_tab: ziggurat_tables::ZigTable,
212 pdf: |f64|: 'static -> f64,
213 zero_case: |&mut R, f64|: 'static -> f64)
214 -> f64 {
215 static SCALE: f64 = (1u64 << 53) as f64;
216 loop {
217 // reimplement the f64 generation as an optimisation suggested
218 // by the Doornik paper: we have a lot of precision-space
219 // (i.e. there are 11 bits of the 64 of a u64 to use after
220 // creating a f64), so we might as well reuse some to save
221 // generating a whole extra random number. (Seems to be 15%
222 // faster.)
223 let bits: u64 = rng.gen();
224 let i = (bits & 0xff) as uint;
225 let f = (bits >> 11) as f64 / SCALE;
226
227 // u is either U(-1, 1) or U(0, 1) depending on if this is a
228 // symmetric distribution or not.
229 let u = if symmetric {2.0 * f - 1.0} else {f};
230 let x = u * x_tab[i];
231
232 let test_x = if symmetric {num::abs(x)} else {x};
233
234 // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
235 if test_x < x_tab[i + 1] {
236 return x;
237 }
238 if i == 0 {
239 return zero_case(rng, u);
240 }
241 // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
242 if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
243 return x;
244 }
245 }
246 }
247
248 #[cfg(test)]
249 mod tests {
250 use {task_rng, Rng, Rand};
251 use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};
252
253 #[deriving(Eq, Show)]
254 struct ConstRand(uint);
255 impl Rand for ConstRand {
256 fn rand<R: Rng>(_: &mut R) -> ConstRand {
257 ConstRand(0)
258 }
259 }
260
261 // 0, 1, 2, 3, ...
262 struct CountingRng { i: u32 }
263 impl Rng for CountingRng {
264 fn next_u32(&mut self) -> u32 {
265 self.i += 1;
266 self.i - 1
267 }
268 fn next_u64(&mut self) -> u64 {
269 self.next_u32() as u64
270 }
271 }
272
273 #[test]
274 fn test_rand_sample() {
275 let mut rand_sample = RandSample::<ConstRand>;
276
277 assert_eq!(rand_sample.sample(&mut task_rng()), ConstRand(0));
278 assert_eq!(rand_sample.ind_sample(&mut task_rng()), ConstRand(0));
279 }
280 #[test]
281 fn test_weighted_choice() {
282 // this makes assumptions about the internal implementation of
283 // WeightedChoice, specifically: it doesn't reorder the items,
284 // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
285 // 1, internally; modulo a modulo operation).
286
287 macro_rules! t (
288 ($items:expr, $expected:expr) => {{
289 let wc = WeightedChoice::new($items);
290 let expected = $expected;
291
292 let mut rng = CountingRng { i: 0 };
293
294 for &val in expected.iter() {
295 assert_eq!(wc.ind_sample(&mut rng), val)
296 }
297 }}
298 );
299
300 t!(vec!(Weighted { weight: 1, item: 10}), [10]);
301
302 // skip some
303 t!(vec!(Weighted { weight: 0, item: 20},
304 Weighted { weight: 2, item: 21},
305 Weighted { weight: 0, item: 22},
306 Weighted { weight: 1, item: 23}),
307 [21,21, 23]);
308
309 // different weights
310 t!(vec!(Weighted { weight: 4, item: 30},
311 Weighted { weight: 3, item: 31}),
312 [30,30,30,30, 31,31,31]);
313
314 // check that we're binary searching
315 // correctly with some vectors of odd
316 // length.
317 t!(vec!(Weighted { weight: 1, item: 40},
318 Weighted { weight: 1, item: 41},
319 Weighted { weight: 1, item: 42},
320 Weighted { weight: 1, item: 43},
321 Weighted { weight: 1, item: 44}),
322 [40, 41, 42, 43, 44]);
323 t!(vec!(Weighted { weight: 1, item: 50},
324 Weighted { weight: 1, item: 51},
325 Weighted { weight: 1, item: 52},
326 Weighted { weight: 1, item: 53},
327 Weighted { weight: 1, item: 54},
328 Weighted { weight: 1, item: 55},
329 Weighted { weight: 1, item: 56}),
330 [50, 51, 52, 53, 54, 55, 56]);
331 }
332
333 #[test] #[should_fail]
334 fn test_weighted_choice_no_items() {
335 WeightedChoice::<int>::new(vec!());
336 }
337 #[test] #[should_fail]
338 fn test_weighted_choice_zero_weight() {
339 WeightedChoice::new(vec!(Weighted { weight: 0, item: 0},
340 Weighted { weight: 0, item: 1}));
341 }
342 #[test] #[should_fail]
343 fn test_weighted_choice_weight_overflows() {
344 let x = (-1) as uint / 2; // x + x + 2 is the overflow
345 WeightedChoice::new(vec!(Weighted { weight: x, item: 0 },
346 Weighted { weight: 1, item: 1 },
347 Weighted { weight: x, item: 2 },
348 Weighted { weight: 1, item: 3 }));
349 }
350 }
librand/distributions/mod.rs:51:68-51:68 -trait- definition:
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
/// Generate a random value.
references:- 1265: impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
66: fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
librand/distributions/range.rs:
66: }
67: impl<Sup: SampleRange> IndependentSample<Sup> for Range<Sup> {
68: fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
librand/distributions/gamma.rs:
145: }
146: impl IndependentSample<f64> for GammaSmallShape {
147: fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
--
308: }
309: impl IndependentSample<f64> for StudentT {
310: fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
librand/distributions/normal.rs:
140: }
141: impl IndependentSample<f64> for LogNormal {
142: fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
librand/distributions/exponential.rs:
83: }
84: impl IndependentSample<f64> for Exp {
85: fn ind_sample<R: Rng>(&self, rng: &mut R) -> f64 {
librand/distributions/mod.rs:
146: impl<T: Clone> IndependentSample<T> for WeightedChoice<T> {
147: fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
librand/distributions/mod.rs:37:69-37:69 -trait- definition:
/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
/// Generate a random value of `Support`, using `rng` as the
references:- 1361: impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
62: fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
librand/distributions/range.rs:
63: impl<Sup: SampleRange> Sample<Sup> for Range<Sup> {
64: #[inline]
librand/distributions/gamma.rs:
215: }
216: impl Sample<f64> for ChiSquared {
217: fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
--
268: }
269: impl Sample<f64> for FisherF {
270: fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
--
305: }
306: impl Sample<f64> for StudentT {
307: fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
librand/distributions/normal.rs:
99: }
100: impl Sample<f64> for Normal {
101: fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
--
137: }
138: impl Sample<f64> for LogNormal {
139: fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
librand/distributions/exponential.rs:
81: impl Sample<f64> for Exp {
82: fn sample<R: Rng>(&mut self, rng: &mut R) -> f64 { self.ind_sample(rng) }
librand/distributions/mod.rs:
142: impl<T: Clone> Sample<T> for WeightedChoice<T> {
143: fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
librand/distributions/mod.rs:102:8-102:8 -struct- definition:
/// ```
pub struct WeightedChoice<T> {
items: Vec<Weighted<T>>,
references:- 5142: impl<T: Clone> Sample<T> for WeightedChoice<T> {
143: fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
--
146: impl<T: Clone> IndependentSample<T> for WeightedChoice<T> {
147: fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
librand/distributions/mod.rs:206:18-206:18 -fn- definition:
fn ziggurat<R:Rng>(
rng: &mut R,
symmetric: bool,
references:- 2librand/distributions/exponential.rs:
46: Exp1(ziggurat(rng, false,
47: &ziggurat_tables::ZIG_EXP_X,
librand/distributions/normal.rs:
60: StandardNormal(ziggurat(
61: rng,
librand/distributions/mod.rs:71:68-71:68 -struct- definition:
/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
/// The numerical weight of this item
references:- 2103: pub struct WeightedChoice<T> {
104: items: Vec<Weighted<T>>,
105: weight_range: Range<uint>
--
114: /// - the total weight is larger than a `uint` can contain.
115: pub fn new(mut items: Vec<Weighted<T>>) -> WeightedChoice<T> {
116: // strictly speaking, this is subsumed by the total weight == 0 case
librand/distributions/mod.rs:58:43-58:43 -struct- definition:
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;
impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
references:- 261: impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
62: fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
--
65: impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
66: fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {