lbf/samplers/
rotation_distr.rs

1use rand::Rng;
2use rand::prelude::Distribution;
3use rand::prelude::IndexedRandom;
4use rand_distr::Normal;
5use rand_distr::Uniform;
6use std::f32::consts::PI;
7
8use jagua_rs::entities::Item;
9use jagua_rs::geometry::geo_enums::RotationRange;
10
11/// Samples a rotation (radians).
12pub trait RotationSampler {
13    fn sample(&self, rng: &mut impl Rng) -> f32;
14}
15
16/// Samples a rotation from a uniform distribution over a given range or a discrete set of rotations.
17pub enum UniformRotDistr {
18    Range(Uniform<f32>),
19    Discrete(Vec<f32>),
20    None,
21}
22
23/// Samples a rotation from a normal distribution over a given range or a discrete set of rotations.
24/// In case of discrete rotations the mean is always returned.
25pub enum NormalRotDistr {
26    Range(Normal<f32>),
27    Discrete(f32),
28}
29
30impl UniformRotDistr {
31    pub fn new(rotations: &RotationRange) -> Self {
32        match rotations {
33            RotationRange::None => UniformRotDistr::None,
34            RotationRange::Continuous => {
35                UniformRotDistr::Range(Uniform::new(0.0, 2.0 * PI).unwrap())
36            }
37            RotationRange::Discrete(a_o) => UniformRotDistr::Discrete(a_o.clone()),
38        }
39    }
40
41    pub fn sample(&self, rng: &mut impl Rng) -> f32 {
42        match self {
43            UniformRotDistr::None => 0.0,
44            UniformRotDistr::Range(u) => u.sample(rng),
45            UniformRotDistr::Discrete(a_o) => *a_o.choose(rng).unwrap(),
46        }
47    }
48}
49
50impl NormalRotDistr {
51    pub fn from_item(item: &Item, r_ref: f32, stddev: f32) -> Self {
52        match item.allowed_orientations().rotations() {
53            RotationRange::None => NormalRotDistr::Discrete(r_ref),
54            RotationRange::Continuous => NormalRotDistr::Range(Normal::new(r_ref, stddev).unwrap()),
55            RotationRange::Discrete(_) => NormalRotDistr::Discrete(r_ref),
56        }
57    }
58
59    pub fn set_mean(&mut self, mean: f32) {
60        match self {
61            NormalRotDistr::Range(n) => {
62                *n = Normal::new(mean, n.std_dev()).unwrap();
63            }
64            NormalRotDistr::Discrete(r) => *r = mean,
65        }
66    }
67
68    pub fn set_stddev(&mut self, stddev: f32) {
69        match self {
70            NormalRotDistr::Range(n) => {
71                *n = Normal::new(n.mean(), stddev).unwrap();
72            }
73            NormalRotDistr::Discrete(_) => {}
74        }
75    }
76
77    pub fn sample(&self, rng: &mut impl Rng) -> f32 {
78        match self {
79            NormalRotDistr::Range(n) => n.sample(rng),
80            NormalRotDistr::Discrete(r) => *r,
81        }
82    }
83}
lbf/samplers/rotation_distr.rs

lbf/samplers/
rotation_distr.rs
