Skip to main content

lbf/samplers/
ls_sampler.rs

1use rand::Rng;
2use rand_distr::Distribution;
3use rand_distr::Normal;
4use std::f32::consts::PI;
5
6use crate::samplers::rotation_distr::NormalRotDistr;
7use jagua_rs::entities::Item;
8use jagua_rs::geometry::DTransformation;
9use jagua_rs::geometry::primitives::Rect;
10
11/// The stddev of translation starts at 1% and ends at 0.05% of the largest dimension of the bounding box.
12pub const SD_TRANSL: (f32, f32) = (0.01, 0.0005);
13
14/// The stddev of rotation starts at 2° and ends at 0.5°.
15pub const SD_ROT: (f32, f32) = (2.0 * PI / 180.0, 0.5 * PI / 180.0);
16
17///Creates `Transformation` samples for a given item.
18///The samples are drawn from normal distributions with decaying standard deviations.
19///Each time an improvement is found, the mean of the distributions is shifted to the new best transformation.
20pub struct LSSampler {
21    normal_x: Normal<f32>,
22    normal_y: Normal<f32>,
23    normal_r: NormalRotDistr,
24    reflected: bool,
25    sd_transl: f32,
26    sd_rot: f32,
27    sd_transl_range: (f32, f32),
28    sd_rot_range: (f32, f32),
29    pub(crate) n_samples: usize,
30}
31
32impl LSSampler {
33    pub fn new(
34        item: &Item,
35        ref_transform: DTransformation,
36        sd_transl_range: (f32, f32),
37        sd_rot_range: (f32, f32),
38    ) -> Self {
39        let sd_transl = sd_transl_range.0;
40        let sd_rot = sd_rot_range.0;
41
42        let normal_x = Normal::new(ref_transform.translation().0, sd_transl).unwrap();
43        let normal_y = Normal::new(ref_transform.translation().1, sd_transl).unwrap();
44        let normal_r = NormalRotDistr::from_item(item, ref_transform.rotation(), sd_rot);
45
46        Self {
47            normal_x,
48            normal_y,
49            normal_r,
50            reflected: ref_transform.reflected,
51            sd_transl,
52            sd_rot,
53            sd_transl_range,
54            sd_rot_range,
55            n_samples: 0,
56        }
57    }
58
59    /// Creates a new sampler with default standard deviation ranges: [SD_TRANSL] and [SD_ROT].
60    pub fn from_defaults(item: &Item, ref_transform: DTransformation, bbox: Rect) -> Self {
61        let max_dim = f32::max(bbox.width(), bbox.height());
62        let sd_transl_range = (SD_TRANSL.0 * max_dim, SD_TRANSL.1 * max_dim);
63        Self::new(item, ref_transform, sd_transl_range, SD_ROT)
64    }
65
66    /// Shifts the mean of the normal distributions to the given reference transformation.
67    pub fn shift_mean(&mut self, ref_transform: DTransformation) {
68        self.normal_x = Normal::new(ref_transform.translation().0, self.sd_transl).unwrap();
69        self.normal_y = Normal::new(ref_transform.translation().1, self.sd_transl).unwrap();
70        self.normal_r.set_mean(ref_transform.rotation());
71        self.reflected = ref_transform.reflected;
72    }
73
74    /// Sets the standard deviation of the normal distributions.
75    pub fn set_stddev(&mut self, stddev_transl: f32, stddev_rot: f32) {
76        assert!(stddev_transl >= 0.0 && stddev_rot >= 0.0);
77
78        self.sd_transl = stddev_transl;
79        self.sd_rot = stddev_rot;
80        self.normal_x = Normal::new(self.normal_x.mean(), self.sd_transl).unwrap();
81        self.normal_y = Normal::new(self.normal_y.mean(), self.sd_transl).unwrap();
82        self.normal_r.set_stddev(self.sd_rot);
83    }
84
85    /// Adjusts the standard deviation according to the fraction of samples that have passed,
86    /// following an exponential decay curve.
87    /// `progress_pct` is a value in [0, 1].
88    ///
89    /// f(0) = init;
90    /// f(1) = end;
91    /// f(x) = init * (end/init)^x;
92    pub fn decay_stddev(&mut self, progress_pct: f32) {
93        let calc_stddev = |(init, end): (f32, f32), pct: f32| init * (end / init).powf(pct);
94        self.set_stddev(
95            calc_stddev(self.sd_transl_range, progress_pct),
96            calc_stddev(self.sd_rot_range, progress_pct),
97        );
98    }
99
100    /// Samples a transformation from the distribution.
101    pub fn sample(&mut self, rng: &mut impl Rng) -> DTransformation {
102        self.n_samples += 1;
103
104        DTransformation::new(
105            self.normal_r.sample(rng),
106            (self.normal_x.sample(rng), self.normal_y.sample(rng)),
107        )
108        .with_reflection(self.reflected)
109    }
110}