rv 0.10.0

Random variables
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

/// Using an enum as a categorical variable.
///
/// In a future where proc_macro is stable, we can offload all of this to a
/// custom derive to have code that looks like this:
///
/// ```
/// #[derive(CategoricalEnum)]
/// enum Color {
///     Red = 0,
///     Blue = 1,
///     Green = 2,
/// }
/// ```
use rv::data::CategoricalDatum;
use rv::prelude::*;

/// We have to assign values 0, ..., n-1 to the enum values so they map to
/// indices in the categorical weights
#[repr(u8)]
#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
enum Color {
    Red = 0,
    Blue = 1,
    Green = 2,
}

/// Then we implement the CatgoricalDatum trait for Color which has methods to
/// convert to and from a usize index.
impl CategoricalDatum for Color {
    fn into_usize(&self) -> usize {
        *self as usize
    }

    fn from_usize(n: usize) -> Self {
        match n {
            0 => Color::Red,
            1 => Color::Blue,
            2 => Color::Green,
            _ => panic!("Cannot convert {} to Color", n),
        }
    }
}

fn main() {
    let mut rng = rand::thread_rng();

    let ctgrl = Categorical::new(&vec![0.25, 0.25, 0.5]).unwrap();

    // Compute the PMF, P(Red).
    let p = ctgrl.pmf(&Color::Red);
    println!("p({:?}) = {}", Color::Red, p);

    // Take 10 draws from {Red, Blue, Green} according to the distribution.
    let xs: Vec<Color> = ctgrl.sample(10, &mut rng);
    println!("xs: {:?}", xs);
}