lbf/samplers/
rotation_distr.rs

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
use rand::Rng;
use rand::prelude::Distribution;
use rand::prelude::IndexedRandom;
use rand_distr::Normal;
use rand_distr::Uniform;

use jagua_rs::entities::item::Item;
use jagua_rs::geometry::geo_enums::AllowedRotation;
use jagua_rs::{PI, fsize};

/// Samples a rotation (radians).
pub trait RotationSampler {
    fn sample(&self, rng: &mut impl Rng) -> fsize;
}

/// Samples a rotation from a uniform distribution over a given range or a discrete set of rotations.
pub enum UniformRotDistr {
    Range(Uniform<fsize>),
    Discrete(Vec<fsize>),
    None,
}

/// Samples a rotation from a normal distribution over a given range or a discrete set of rotations.
/// In case of discrete rotations the mean is always returned.
pub enum NormalRotDistr {
    Range(Normal<fsize>),
    Discrete(fsize),
    None,
}

impl UniformRotDistr {
    pub fn from_item(item: &Item) -> Self {
        match &item.allowed_rotation {
            AllowedRotation::None => UniformRotDistr::None,
            AllowedRotation::Continuous => {
                UniformRotDistr::Range(Uniform::new(0.0, 2.0 * PI).unwrap())
            }
            AllowedRotation::Discrete(a_o) => UniformRotDistr::Discrete(a_o.clone()),
        }
    }

    pub fn sample(&self, rng: &mut impl Rng) -> fsize {
        match self {
            UniformRotDistr::None => 0.0,
            UniformRotDistr::Range(u) => u.sample(rng),
            UniformRotDistr::Discrete(a_o) => *a_o.choose(rng).unwrap(),
        }
    }
}

impl NormalRotDistr {
    pub fn from_item(item: &Item, r_ref: fsize, stddev: fsize) -> Self {
        match &item.allowed_rotation {
            AllowedRotation::None => NormalRotDistr::None,
            AllowedRotation::Continuous => {
                NormalRotDistr::Range(Normal::new(r_ref, stddev).unwrap())
            }
            AllowedRotation::Discrete(_) => NormalRotDistr::Discrete(r_ref),
        }
    }

    pub fn set_mean(&mut self, mean: fsize) {
        match self {
            NormalRotDistr::Range(n) => {
                *n = Normal::new(mean, n.std_dev()).unwrap();
            }
            NormalRotDistr::Discrete(_) | NormalRotDistr::None => {}
        }
    }

    pub fn set_stddev(&mut self, stddev: fsize) {
        match self {
            NormalRotDistr::Range(n) => {
                *n = Normal::new(n.mean(), stddev).unwrap();
            }
            NormalRotDistr::Discrete(_) | NormalRotDistr::None => {}
        }
    }

    pub fn sample(&self, rng: &mut impl Rng) -> fsize {
        match self {
            NormalRotDistr::None => 0.0,
            NormalRotDistr::Range(n) => n.sample(rng),
            NormalRotDistr::Discrete(r) => *r,
        }
    }
}