fraug/augmenters/
addnoise.rs

1use super::base::Augmenter;
2use rand::{distr::Uniform, prelude::*, random_range};
3use rand_distr::Normal;
4use  tracing::{info_span};
5
6/// Augmenter that allows different types of noise injection
7///
8/// Noise types:
9/// - Uniform: Adds uniform noise within the given bounds given through the parameter `bounds`
10/// - Gaussian: Adds gaussian noise with the specified mean and standard deviation according to the corresponding parameters
11/// - Spike: Adds a spike in the series with a random magnitude (in the range specified by `bounds` of the standard deviation of the original time series
12/// - Slope: Adds a linear slope trend to the series with a random slope in the range specified by `bounds`
13pub 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
22/// Enum to specify the noise type for the AddNoise augmenter
23pub 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                // Calculate std dev of x
74                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                // Add spike in random location with random magnitude
79                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}