1use crate::NeuralResult;
8use scirs2_core::ndarray::{s, Array1, Array2, Array3};
9use scirs2_core::random::essentials::{Normal, Uniform, Uniform as RandUniform};
10use scirs2_core::random::{Distribution, RngExt};
11use scirs2_core::{ChaCha8Rng, SeedableRng};
12use sklears_core::types::FloatBounds;
13use std::collections::HashMap;
14
15pub struct AugmentationPipeline<T: FloatBounds> {
17 transformations: Vec<Box<dyn Transformation<T>>>,
18 probability: T,
19 seed: Option<u64>,
20 rng: ChaCha8Rng,
21}
22
23impl<T: FloatBounds> AugmentationPipeline<T> {
24 pub fn new() -> Self {
26 Self {
27 transformations: Vec::new(),
28 probability: T::one(),
29 seed: None,
30 rng: ChaCha8Rng::seed_from_u64(42),
31 }
32 }
33
34 pub fn with_seed(seed: u64) -> Self {
36 Self {
37 transformations: Vec::new(),
38 probability: T::one(),
39 seed: Some(seed),
40 rng: ChaCha8Rng::seed_from_u64(seed),
41 }
42 }
43
44 pub fn probability(mut self, prob: T) -> Self {
46 self.probability = prob;
47 self
48 }
49
50 pub fn add_transformation(mut self, transform: Box<dyn Transformation<T>>) -> Self {
52 self.transformations.push(transform);
53 self
54 }
55
56 pub fn apply(&mut self, data: &Array2<T>) -> NeuralResult<Array2<T>> {
58 let apply_prob: f64 = self.probability.to_f64().unwrap_or(1.0);
60 if self.rng.random::<f64>() > apply_prob {
61 return Ok(data.clone());
62 }
63
64 let mut result = data.clone();
65 for transform in &mut self.transformations {
66 result = transform.apply(&result, &mut self.rng)?;
67 }
68 Ok(result)
69 }
70
71 pub fn apply_batch(&mut self, batch: &Array3<T>) -> NeuralResult<Array3<T>> {
73 let (batch_size, height, width) = batch.dim();
74 let mut result = Array3::zeros((batch_size, height, width));
75
76 for i in 0..batch_size {
77 let sample = batch.slice(s![i, .., ..]).to_owned();
78 let augmented = self.apply(&sample)?;
79 result.slice_mut(s![i, .., ..]).assign(&augmented);
80 }
81
82 Ok(result)
83 }
84
85 pub fn reseed(&mut self, seed: u64) {
87 self.seed = Some(seed);
88 self.rng = ChaCha8Rng::seed_from_u64(seed);
89 }
90}
91
92impl<T: FloatBounds> Default for AugmentationPipeline<T> {
93 fn default() -> Self {
94 Self::new()
95 }
96}
97
98pub trait Transformation<T: FloatBounds> {
100 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>>;
102
103 fn name(&self) -> &str;
105
106 fn parameters(&self) -> HashMap<String, String>;
108}
109
110#[derive(Debug)]
112pub struct GaussianNoise<T: FloatBounds> {
113 mean: T,
114 std: T,
115 probability: T,
116 name: String,
117}
118
119impl<T: FloatBounds> GaussianNoise<T> {
120 pub fn new(mean: T, std: T) -> Self {
122 Self {
123 mean,
124 std,
125 probability: T::one(),
126 name: "GaussianNoise".to_string(),
127 }
128 }
129
130 pub fn with_probability(mut self, prob: T) -> Self {
132 self.probability = prob;
133 self
134 }
135}
136
137impl<T: FloatBounds> Transformation<T> for GaussianNoise<T> {
138 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>> {
139 let prob: f64 = self.probability.to_f64().unwrap_or(1.0);
140 if rng.random::<f64>() > prob {
141 return Ok(data.clone());
142 }
143
144 let mean: f64 = self.mean.to_f64().unwrap_or(0.0);
145 let std: f64 = self.std.to_f64().unwrap_or(1.0);
146
147 let normal = Normal::new(mean, std).map_err(|_| {
148 sklears_core::error::SklearsError::InvalidParameter {
149 name: "noise_parameters".to_string(),
150 reason: "Invalid normal distribution parameters".to_string(),
151 }
152 })?;
153
154 let mut result = data.clone();
155 result.mapv_inplace(|x| {
156 let noise = T::from(normal.sample(rng)).unwrap_or_else(T::zero);
157 x + noise
158 });
159
160 Ok(result)
161 }
162
163 fn name(&self) -> &str {
164 &self.name
165 }
166
167 fn parameters(&self) -> HashMap<String, String> {
168 let mut params = HashMap::new();
169 params.insert("mean".to_string(), format!("{:?}", self.mean));
170 params.insert("std".to_string(), format!("{:?}", self.std));
171 params.insert("probability".to_string(), format!("{:?}", self.probability));
172 params
173 }
174}
175
176#[derive(Debug)]
178pub struct UniformNoise<T: FloatBounds> {
179 low: T,
180 high: T,
181 probability: T,
182 name: String,
183}
184
185impl<T: FloatBounds> UniformNoise<T> {
186 pub fn new(low: T, high: T) -> Self {
188 Self {
189 low,
190 high,
191 probability: T::one(),
192 name: "UniformNoise".to_string(),
193 }
194 }
195
196 pub fn with_probability(mut self, prob: T) -> Self {
198 self.probability = prob;
199 self
200 }
201}
202
203impl<T: FloatBounds> Transformation<T> for UniformNoise<T> {
204 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>> {
205 let prob: f64 = self.probability.to_f64().unwrap_or(1.0);
206 if rng.random::<f64>() > prob {
207 return Ok(data.clone());
208 }
209
210 let low: f64 = self.low.to_f64().unwrap_or(0.0);
211 let high: f64 = self.high.to_f64().unwrap_or(1.0);
212
213 let uniform = Uniform::new(low, high).expect("valid distribution params");
214 let mut result = data.clone();
215
216 result.mapv_inplace(|x| {
217 let noise = T::from(uniform.sample(rng)).unwrap_or_else(T::zero);
218 x + noise
219 });
220
221 Ok(result)
222 }
223
224 fn name(&self) -> &str {
225 &self.name
226 }
227
228 fn parameters(&self) -> HashMap<String, String> {
229 let mut params = HashMap::new();
230 params.insert("low".to_string(), format!("{:?}", self.low));
231 params.insert("high".to_string(), format!("{:?}", self.high));
232 params.insert("probability".to_string(), format!("{:?}", self.probability));
233 params
234 }
235}
236
237#[derive(Debug)]
239pub struct FeatureDropout<T: FloatBounds> {
240 dropout_rate: T,
241 probability: T,
242 name: String,
243}
244
245impl<T: FloatBounds> FeatureDropout<T> {
246 pub fn new(dropout_rate: T) -> Self {
248 Self {
249 dropout_rate,
250 probability: T::one(),
251 name: "FeatureDropout".to_string(),
252 }
253 }
254
255 pub fn with_probability(mut self, prob: T) -> Self {
257 self.probability = prob;
258 self
259 }
260}
261
262impl<T: FloatBounds> Transformation<T> for FeatureDropout<T> {
263 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>> {
264 let prob: f64 = self.probability.to_f64().unwrap_or(1.0);
265 if rng.random::<f64>() > prob {
266 return Ok(data.clone());
267 }
268
269 let dropout_rate: f64 = self.dropout_rate.to_f64().unwrap_or(0.0);
270 let mut result = data.clone();
271
272 for mut column in result.columns_mut() {
274 if rng.random::<f64>() < dropout_rate {
275 column.fill(T::zero());
276 }
277 }
278
279 Ok(result)
280 }
281
282 fn name(&self) -> &str {
283 &self.name
284 }
285
286 fn parameters(&self) -> HashMap<String, String> {
287 let mut params = HashMap::new();
288 params.insert(
289 "dropout_rate".to_string(),
290 format!("{:?}", self.dropout_rate),
291 );
292 params.insert("probability".to_string(), format!("{:?}", self.probability));
293 params
294 }
295}
296
297#[derive(Debug)]
299pub struct FeatureScaling<T: FloatBounds> {
300 scale_range: (T, T),
301 probability: T,
302 name: String,
303}
304
305impl<T: FloatBounds> FeatureScaling<T> {
306 pub fn new(min_scale: T, max_scale: T) -> Self {
308 Self {
309 scale_range: (min_scale, max_scale),
310 probability: T::one(),
311 name: "FeatureScaling".to_string(),
312 }
313 }
314
315 pub fn with_probability(mut self, prob: T) -> Self {
317 self.probability = prob;
318 self
319 }
320}
321
322impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> Transformation<T> for FeatureScaling<T> {
323 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>> {
324 let prob: f64 = self.probability.to_f64().unwrap_or(1.0);
325 if rng.random::<f64>() > prob {
326 return Ok(data.clone());
327 }
328
329 let min_scale: f64 = self.scale_range.0.to_f64().unwrap_or(0.8);
330 let max_scale: f64 = self.scale_range.1.to_f64().unwrap_or(1.2);
331
332 let uniform = RandUniform::new(min_scale, max_scale).expect("valid distribution params");
333 let scale_factor = T::from(uniform.sample(rng)).unwrap_or_else(T::one);
334
335 Ok(data * scale_factor)
336 }
337
338 fn name(&self) -> &str {
339 &self.name
340 }
341
342 fn parameters(&self) -> HashMap<String, String> {
343 let mut params = HashMap::new();
344 params.insert("min_scale".to_string(), format!("{:?}", self.scale_range.0));
345 params.insert("max_scale".to_string(), format!("{:?}", self.scale_range.1));
346 params.insert("probability".to_string(), format!("{:?}", self.probability));
347 params
348 }
349}
350
351#[derive(Debug)]
353pub struct FeaturePermutation<T: FloatBounds> {
354 permutation_ratio: T,
355 probability: T,
356 name: String,
357}
358
359impl<T: FloatBounds> FeaturePermutation<T> {
360 pub fn new(permutation_ratio: T) -> Self {
362 Self {
363 permutation_ratio,
364 probability: T::one(),
365 name: "FeaturePermutation".to_string(),
366 }
367 }
368
369 pub fn with_probability(mut self, prob: T) -> Self {
371 self.probability = prob;
372 self
373 }
374}
375
376impl<T: FloatBounds> Transformation<T> for FeaturePermutation<T> {
377 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>> {
378 let prob: f64 = self.probability.to_f64().unwrap_or(1.0);
379 if rng.random::<f64>() > prob {
380 return Ok(data.clone());
381 }
382
383 let ratio: f64 = self.permutation_ratio.to_f64().unwrap_or(0.1);
384 let (n_samples, n_features) = data.dim();
385 let n_permute = ((n_features as f64) * ratio) as usize;
386
387 let mut result = data.clone();
388
389 let mut features_to_permute = Vec::new();
391 for _ in 0..n_permute {
392 features_to_permute.push(rng.random_range(0..n_features));
393 }
394
395 for &feature_idx in &features_to_permute {
397 let mut column = result.column(feature_idx).to_owned();
398
399 for i in (1..n_samples).rev() {
401 let j = rng.random_range(0..=i);
402 column.swap(i, j);
403 }
404
405 result.column_mut(feature_idx).assign(&column);
406 }
407
408 Ok(result)
409 }
410
411 fn name(&self) -> &str {
412 &self.name
413 }
414
415 fn parameters(&self) -> HashMap<String, String> {
416 let mut params = HashMap::new();
417 params.insert(
418 "permutation_ratio".to_string(),
419 format!("{:?}", self.permutation_ratio),
420 );
421 params.insert("probability".to_string(), format!("{:?}", self.probability));
422 params
423 }
424}
425
426#[derive(Debug)]
428pub struct TimeSeriesAugmentation<T: FloatBounds> {
429 transformations: Vec<TimeSeriesTransform>,
430 probability: T,
431 name: String,
432}
433
434#[derive(Debug, Clone)]
436pub enum TimeSeriesTransform {
437 TimeWarp {
439 sigma: f64,
441 },
442 MagnitudeWarp {
444 sigma: f64,
446 },
447 WindowSlicing {
449 ratio: f64,
451 },
452 Jittering {
454 sigma: f64,
456 },
457}
458
459impl<T: FloatBounds> TimeSeriesAugmentation<T> {
460 pub fn new() -> Self {
462 Self {
463 transformations: Vec::new(),
464 probability: T::one(),
465 name: "TimeSeriesAugmentation".to_string(),
466 }
467 }
468
469 pub fn add_transform(mut self, transform: TimeSeriesTransform) -> Self {
471 self.transformations.push(transform);
472 self
473 }
474
475 pub fn with_probability(mut self, prob: T) -> Self {
477 self.probability = prob;
478 self
479 }
480}
481
482impl<T: FloatBounds> Transformation<T> for TimeSeriesAugmentation<T> {
483 fn apply(&mut self, data: &Array2<T>, rng: &mut ChaCha8Rng) -> NeuralResult<Array2<T>> {
484 let prob: f64 = self.probability.to_f64().unwrap_or(1.0);
485 if rng.random::<f64>() > prob {
486 return Ok(data.clone());
487 }
488
489 let mut result = data.clone();
490
491 for transform in &self.transformations {
492 result = match transform {
493 TimeSeriesTransform::Jittering { sigma } => {
494 self.apply_jittering(&result, *sigma, rng)?
495 }
496 TimeSeriesTransform::TimeWarp { sigma } => {
497 self.apply_time_warp(&result, *sigma, rng)?
498 }
499 TimeSeriesTransform::MagnitudeWarp { sigma } => {
500 self.apply_magnitude_warp(&result, *sigma, rng)?
501 }
502 TimeSeriesTransform::WindowSlicing { ratio } => {
503 self.apply_window_slicing(&result, *ratio, rng)?
504 }
505 };
506 }
507
508 Ok(result)
509 }
510
511 fn name(&self) -> &str {
512 &self.name
513 }
514
515 fn parameters(&self) -> HashMap<String, String> {
516 let mut params = HashMap::new();
517 params.insert(
518 "num_transforms".to_string(),
519 self.transformations.len().to_string(),
520 );
521 params.insert("probability".to_string(), format!("{:?}", self.probability));
522 for (i, transform) in self.transformations.iter().enumerate() {
523 params.insert(format!("transform_{}", i), format!("{:?}", transform));
524 }
525 params
526 }
527}
528
529impl<T: FloatBounds> TimeSeriesAugmentation<T> {
530 fn apply_jittering(
531 &self,
532 data: &Array2<T>,
533 sigma: f64,
534 rng: &mut ChaCha8Rng,
535 ) -> NeuralResult<Array2<T>> {
536 let normal = Normal::new(0.0, sigma).map_err(|_| {
537 sklears_core::error::SklearsError::InvalidParameter {
538 name: "jittering_sigma".to_string(),
539 reason: "Invalid sigma for jittering".to_string(),
540 }
541 })?;
542
543 let mut result = data.clone();
544 result.mapv_inplace(|x| {
545 let noise = T::from(normal.sample(rng)).unwrap_or_else(T::zero);
546 x + noise
547 });
548
549 Ok(result)
550 }
551
552 fn apply_time_warp(
553 &self,
554 data: &Array2<T>,
555 sigma: f64,
556 rng: &mut ChaCha8Rng,
557 ) -> NeuralResult<Array2<T>> {
558 let (n_samples, n_features) = data.dim();
559 let mut result = Array2::zeros((n_samples, n_features));
560
561 let normal = Normal::new(0.0, sigma).map_err(|_| {
563 sklears_core::error::SklearsError::InvalidParameter {
564 name: "time_warp_sigma".to_string(),
565 reason: "Invalid sigma for time warping".to_string(),
566 }
567 })?;
568
569 for i in 0..n_samples {
570 let original_row = data.row(i);
571 let mut warped_row = Array1::zeros(n_features);
572
573 for j in 0..n_features {
574 let warp_factor = 1.0 + normal.sample(rng);
576 let warped_idx = (j as f64 * warp_factor) as usize;
577
578 if warped_idx < n_features {
579 warped_row[j] = original_row[warped_idx];
580 } else {
581 warped_row[j] = original_row[n_features - 1];
582 }
583 }
584
585 result.row_mut(i).assign(&warped_row);
586 }
587
588 Ok(result)
589 }
590
591 fn apply_magnitude_warp(
592 &self,
593 data: &Array2<T>,
594 sigma: f64,
595 rng: &mut ChaCha8Rng,
596 ) -> NeuralResult<Array2<T>> {
597 let normal = Normal::new(1.0, sigma).map_err(|_| {
598 sklears_core::error::SklearsError::InvalidParameter {
599 name: "magnitude_warp_sigma".to_string(),
600 reason: "Invalid sigma for magnitude warping".to_string(),
601 }
602 })?;
603
604 let mut result = data.clone();
605
606 let (n_samples, n_features) = data.dim();
608 for i in 0..n_samples {
609 for j in 0..n_features {
610 let scale_factor = T::from(normal.sample(rng)).unwrap_or_else(T::one);
611 result[[i, j]] *= scale_factor;
612 }
613 }
614
615 Ok(result)
616 }
617
618 fn apply_window_slicing(
619 &self,
620 data: &Array2<T>,
621 ratio: f64,
622 rng: &mut ChaCha8Rng,
623 ) -> NeuralResult<Array2<T>> {
624 let (n_samples, n_features) = data.dim();
625 let window_size = ((n_features as f64) * ratio) as usize;
626
627 if window_size == 0 || window_size >= n_features {
628 return Ok(data.clone());
629 }
630
631 let mut result = Array2::zeros((n_samples, n_features));
632
633 for i in 0..n_samples {
634 let start_idx = rng.random_range(0..(n_features - window_size + 1));
635 let original_row = data.row(i);
636
637 for j in 0..n_features {
639 if j < window_size {
640 result[[i, j]] = original_row[start_idx + j];
641 } else {
642 result[[i, j]] = original_row[start_idx + window_size - 1];
644 }
645 }
646 }
647
648 Ok(result)
649 }
650}
651
652impl<T: FloatBounds> Default for TimeSeriesAugmentation<T> {
653 fn default() -> Self {
654 Self::new()
655 }
656}
657
658pub struct AugmentationBuilder<T: FloatBounds> {
660 pipeline: AugmentationPipeline<T>,
661}
662
663impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> AugmentationBuilder<T> {
664 pub fn new() -> Self {
666 Self {
667 pipeline: AugmentationPipeline::new(),
668 }
669 }
670
671 pub fn with_seed(seed: u64) -> Self {
673 Self {
674 pipeline: AugmentationPipeline::with_seed(seed),
675 }
676 }
677
678 pub fn probability(mut self, prob: T) -> Self {
680 self.pipeline = self.pipeline.probability(prob);
681 self
682 }
683
684 pub fn gaussian_noise(mut self, mean: T, std: T) -> Self {
686 let noise = GaussianNoise::new(mean, std);
687 self.pipeline = self.pipeline.add_transformation(Box::new(noise));
688 self
689 }
690
691 pub fn uniform_noise(mut self, low: T, high: T) -> Self {
693 let noise = UniformNoise::new(low, high);
694 self.pipeline = self.pipeline.add_transformation(Box::new(noise));
695 self
696 }
697
698 pub fn feature_dropout(mut self, rate: T) -> Self {
700 let dropout = FeatureDropout::new(rate);
701 self.pipeline = self.pipeline.add_transformation(Box::new(dropout));
702 self
703 }
704
705 pub fn feature_scaling(mut self, min_scale: T, max_scale: T) -> Self {
707 let scaling = FeatureScaling::new(min_scale, max_scale);
708 self.pipeline = self.pipeline.add_transformation(Box::new(scaling));
709 self
710 }
711
712 pub fn feature_permutation(mut self, ratio: T) -> Self {
714 let permutation = FeaturePermutation::new(ratio);
715 self.pipeline = self.pipeline.add_transformation(Box::new(permutation));
716 self
717 }
718
719 pub fn time_series_jittering(mut self, sigma: f64) -> Self {
721 let ts_aug =
722 TimeSeriesAugmentation::new().add_transform(TimeSeriesTransform::Jittering { sigma });
723 self.pipeline = self.pipeline.add_transformation(Box::new(ts_aug));
724 self
725 }
726
727 pub fn build(self) -> AugmentationPipeline<T> {
729 self.pipeline
730 }
731}
732
733impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> Default for AugmentationBuilder<T> {
734 fn default() -> Self {
735 Self::new()
736 }
737}
738
739#[allow(non_snake_case)]
740#[cfg(test)]
741mod tests {
742 use super::*;
743
744 #[test]
745 fn test_gaussian_noise() -> NeuralResult<()> {
746 let mut noise = GaussianNoise::new(0.0f32, 0.1);
747 let data = Array2::ones((5, 10));
748 let mut rng = ChaCha8Rng::seed_from_u64(42);
749
750 let result = noise.apply(&data, &mut rng)?;
751 assert_eq!(result.shape(), data.shape());
752
753 assert_ne!(result, data);
755
756 Ok(())
757 }
758
759 #[test]
760 fn test_feature_dropout() -> NeuralResult<()> {
761 let mut dropout = FeatureDropout::new(0.5f32);
762 let data = Array2::ones((5, 10));
763 let mut rng = ChaCha8Rng::seed_from_u64(42);
764
765 let result = dropout.apply(&data, &mut rng)?;
766 assert_eq!(result.shape(), data.shape());
767
768 let zero_features = result
770 .columns()
771 .into_iter()
772 .filter(|col| col.iter().all(|&x| x == 0.0))
773 .count();
774
775 assert!(zero_features > 0, "Expected some features to be dropped");
776
777 Ok(())
778 }
779
780 #[test]
781 fn test_augmentation_pipeline() -> NeuralResult<()> {
782 let mut pipeline = AugmentationBuilder::with_seed(42)
783 .gaussian_noise(0.0, 0.1)
784 .feature_dropout(0.2)
785 .feature_scaling(0.9, 1.1)
786 .build();
787
788 let data = Array2::ones((5, 10));
789 let result = pipeline.apply(&data)?;
790
791 assert_eq!(result.shape(), data.shape());
792 assert_ne!(result, data);
793
794 Ok(())
795 }
796
797 #[test]
798 fn test_time_series_augmentation() -> NeuralResult<()> {
799 let mut ts_aug = TimeSeriesAugmentation::new()
800 .add_transform(TimeSeriesTransform::Jittering { sigma: 0.1 })
801 .add_transform(TimeSeriesTransform::MagnitudeWarp { sigma: 0.1 });
802
803 let data = Array2::from_shape_fn((3, 20), |(i, j)| (i * 20 + j) as f32);
804 let mut rng = ChaCha8Rng::seed_from_u64(42);
805
806 let result = ts_aug.apply(&data, &mut rng)?;
807 assert_eq!(result.shape(), data.shape());
808
809 Ok(())
810 }
811
812 #[test]
813 fn test_transformation_parameters() {
814 let noise = GaussianNoise::new(0.0f32, 1.0);
815 let params = noise.parameters();
816
817 assert!(params.contains_key("mean"));
818 assert!(params.contains_key("std"));
819 assert!(params.contains_key("probability"));
820 assert_eq!(noise.name(), "GaussianNoise");
821 }
822}