fraug/augmenters/
addnoise.rs1use super::base::Augmenter;
2use rand::{distr::Uniform, prelude::*, random_range};
3use rand_distr::Normal;
4use tracing::{info_span};
5
6pub struct AddNoise {
14 pub name: String,
15 pub noise_type: NoiseType,
16 pub bounds: Option<(f64, f64)>,
17 pub mean: Option<f64>,
18 pub std_dev: Option<f64>,
19 pub p: f64,
20}
21
22pub enum NoiseType {
24 Uniform,
25 Gaussian,
26 Spike,
27 Slope,
28}
29
30impl AddNoise {
31 pub fn new(
32 noise_type: NoiseType,
33 bounds: Option<(f64, f64)>,
34 mean: Option<f64>,
35 std_dev: Option<f64>,
36 ) -> Self {
37 AddNoise {
38 name: "AddNoise".to_string(),
39 noise_type: noise_type,
40 bounds: bounds,
41 mean: mean,
42 std_dev: std_dev,
43 p: 1.0,
44 }
45 }
46}
47
48impl Augmenter for AddNoise {
49 fn augment_one(&self, x: &[f64]) -> Vec<f64> {
50 let span = info_span!("", step = "augment_one");
51 let _enter = span.enter();
52 match self.noise_type {
53 NoiseType::Uniform => {
54 let bounds = self.bounds.expect("Bounds not specified");
55
56 let mut rng = rand::rng();
57 let dist = Uniform::new(bounds.0, bounds.1)
58 .expect("Couldn't create uniform distribution from specified bounds");
59 x.iter().map(|val| *val + dist.sample(&mut rng)).collect()
60 }
61 NoiseType::Gaussian => {
62 let mean = self.mean.expect("Mean not specified");
63 let std_dev = self.std_dev.expect("Standard deviation not specified");
64
65 let mut rng = rand::rng();
66 let dist = Normal::new(mean, std_dev)
67 .expect("Couldn't create normal distribution from specified mean and standard deviation");
68 x.iter().map(|val| *val + dist.sample(&mut rng)).collect()
69 }
70 NoiseType::Spike => {
71 let bounds = self.bounds.expect("Bounds not specified");
72
73 let n = x.len() as f64;
75 let mean = x.iter().sum::<f64>() / n;
76 let std_dev = (x.iter().map(|&val| (val - mean).powi(2)).sum::<f64>() / n).sqrt();
77
78 let idx: usize = random_range(0..n as usize);
80 let magnitude: f64 = random_range(bounds.0..bounds.1);
81
82 let mut res = x.to_vec();
83 res[idx] = magnitude * std_dev;
84 res
85 }
86 NoiseType::Slope => {
87 let bounds = self.bounds.expect("Bounds not specified");
88
89 let slope: f64 = random_range(bounds.0..bounds.1);
90 x.iter()
91 .enumerate()
92 .map(|(i, val)| *val + i as f64 * slope)
93 .collect()
94 }
95 }
96 }
97
98 fn get_probability(&self) -> f64 {
99 self.p
100 }
101
102 fn set_probability(&mut self, probability: f64) {
103 self.p = probability;
104 }
105
106 fn get_name(&self) ->String {
107 self.name.clone()
108 }
109
110}