src/librand/distributions/mod.rs

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