fraug/augmenters/
rotation.rs

1use super::base::Augmenter;
2use tracing::{info_span};
3/// Augmenter that rotates the data 180 degrees around `anchor`
4pub struct Rotation {
5    pub name: String,
6    pub anchor: f64,
7    p: f64,
8}
9
10impl Rotation {
11    pub fn new(anchor: f64) -> Self {
12        Rotation {
13            name: "Rotation".to_string(),
14            anchor,
15            p: 1.0,
16        }
17    }
18}
19
20impl Augmenter for Rotation {
21    fn augment_one(&self, x: &[f64]) -> Vec<f64> {
22        let span = info_span!("", step = "augment_one");
23        let _enter = span.enter();
24
25        x.iter()
26            .map(|val| (*val - self.anchor) * -1.0 + self.anchor)
27            .collect()
28    }
29
30    fn get_probability(&self) -> f64 {
31        self.p
32    }
33
34    fn set_probability(&mut self, probability: f64) {
35        self.p = probability;
36    }
37
38    fn get_name(&self) ->String {
39        self.name.clone()
40    }
41}