fraug/transforms/
fastfourier.rs

1use crate::Dataset;
2use rustfft::{FftPlanner, num_complex::Complex};
3
4use rayon::prelude::*;
5
6/// Converts each real-valued time series in the dataset into its frequency domain representation,
7/// storing the result as interleaved real and imaginary parts: [re0, im0, re1, im1, ...]
8pub fn dataset_fft(dataset: &Dataset, parallel: bool) -> Dataset {
9    let freq_features: Vec<Vec<f64>> = if parallel {
10        dataset
11            .features
12            .par_iter()
13            .map(|sample| {
14                let len = sample.len();
15                let mut planner = FftPlanner::new();
16                let fft = planner.plan_fft_forward(len);
17                let mut buffer: Vec<Complex<f64>> =
18                    sample.iter().map(|&x| Complex { re: x, im: 0.0 }).collect();
19                fft.process(&mut buffer);
20                let mut spectrum = Vec::with_capacity(2 * len);
21                for c in buffer {
22                    spectrum.push(c.re);
23                    spectrum.push(c.im);
24                }
25                spectrum
26            })
27            .collect()
28    } else {
29        dataset
30            .features
31            .iter()
32            .map(|sample| {
33                let len = sample.len();
34                let mut planner = FftPlanner::new();
35                let fft = planner.plan_fft_forward(len);
36                let mut buffer: Vec<Complex<f64>> =
37                    sample.iter().map(|&x| Complex { re: x, im: 0.0 }).collect();
38                fft.process(&mut buffer);
39                let mut spectrum = Vec::with_capacity(2 * len);
40                for c in buffer {
41                    spectrum.push(c.re);
42                    spectrum.push(c.im);
43                }
44                spectrum
45            })
46            .collect()
47    };
48
49    Dataset {
50        features: freq_features,
51        labels: dataset.labels.clone(),
52    }
53}
54
55
56/// Reconstructs each time series from its frequency domain representation (interleaved real/imag parts),
57pub fn dataset_ifft(dataset: &Dataset, parallel: bool) -> Dataset {
58    let time_features: Vec<Vec<f64>> = if parallel {
59        dataset
60            .features
61            .par_iter()
62            .map(|sample| {
63                let len = sample.len() / 2;
64                let mut planner = FftPlanner::new();
65                let ifft = planner.plan_fft_inverse(len);
66                let mut buffer: Vec<Complex<f64>> = (0..len)
67                    .map(|i| Complex {
68                        re: sample[2 * i],
69                        im: sample[2 * i + 1],
70                    })
71                    .collect();
72                ifft.process(&mut buffer);
73                buffer.iter().map(|c| c.re / len as f64).collect()
74            })
75            .collect()
76    } else {
77        dataset
78            .features
79            .iter()
80            .map(|sample| {
81                let len = sample.len() / 2;
82                let mut planner = FftPlanner::new();
83                let ifft = planner.plan_fft_inverse(len);
84                let mut buffer: Vec<Complex<f64>> = (0..len)
85                    .map(|i| Complex {
86                        re: sample[2 * i],
87                        im: sample[2 * i + 1],
88                    })
89                    .collect();
90                ifft.process(&mut buffer);
91                buffer.iter().map(|c| c.re / len as f64).collect()
92            })
93            .collect()
94    };
95
96    Dataset {
97        features: time_features,
98        labels: dataset.labels.clone(),
99    }
100}