fft-rs/src/lib.rs
src/lib.rsRaw
1use num::complex::{Complex, Complex64};
2use std::f64::consts::PI;
3
4// fft computes the Fast Fourier Transform
5pub fn fft(x: &Vec<f64>) -> Vec<Complex64> {
6 let mut x_compl: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); x.len()];
7 for i in 0..x.len() {
8 x_compl[i] = Complex::new(x[i], 0_f64);
9 }
10 fft_compl(x_compl)
11}
12
13fn fft_compl(x: Vec<Complex64>) -> Vec<Complex64> {
14 let N = x.len();
15 if N % 2 > 0 {
16 panic!("not a power of 2");
17 } else if N <= 2 {
18 return dft_compl(x);
19 }
20
21 let mut x_even: Vec<Complex64> = Vec::new();
22 let mut x_odd: Vec<Complex64> = Vec::new();
23 for i in 0..x.len() {
24 if i % 2 == 0 {
25 x_even.push(x[i]);
26 } else {
27 x_odd.push(x[i]);
28 }
29 }
30 let mut x_even_cmplx = fft_compl(x_even);
31 let mut x_odd_cmplx = fft_compl(x_odd);
32
33 let mut w = Complex::new(0_f64, 2_f64 * PI / N as f64);
34 let mut f_k: Vec<Complex64> = Vec::new();
35 for k in 0..x.len() {
36 let k_compl = Complex::new(k as f64, 0_f64);
37 f_k.push((w * k_compl).exp());
38 }
39
40 let mut r: Vec<Complex64> = Vec::new();
41 let mut aa = add_vv(
42 x_even_cmplx.clone(),
43 mul_vv_el(x_odd_cmplx.clone(), f_k.clone()[0..x.len() / 2].to_vec()),
44 );
45 let mut bb = add_vv(
46 x_even_cmplx.clone(),
47 mul_vv_el(x_odd_cmplx.clone(), f_k.clone()[x.len() / 2..].to_vec()),
48 );
49 r.append(&mut aa);
50 r.append(&mut bb);
51
52 r
53}
54
55// ifft computes the Inverse Fast Fourier Transform
56pub fn ifft(x: &Vec<Complex64>) -> Vec<f64> {
57 // use the IFFT method of computing conjugates, then FFT, then conjugate again, and then divide
58 // by N
59 let mut x_conj: Vec<Complex64> = Vec::new();
60 for i in 0..x.len() {
61 x_conj.push(x[i].conj());
62 }
63
64 let x_res = fft_compl(x_conj);
65
66 let mut r: Vec<Complex64> = Vec::new();
67 for i in 0..x_res.len() {
68 r.push(x_res[i].conj());
69 }
70
71 let n = x.len() as f64;
72 let mut rr: Vec<f64> = Vec::new();
73 for i in 0..r.len() {
74 r[i] = r[i] / Complex::new(n, 0_f64);
75 rr.push(r[i].re);
76 }
77 rr
78}
79
80// dft computes the Discrete Fourier Transform
81pub fn dft(x: &Vec<f64>) -> Vec<Complex64> {
82 let mut x_compl: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); x.len()];
83 for i in 0..x.len() {
84 x_compl[i] = Complex::new(x[i], 0_f64);
85 }
86 dft_compl(x_compl)
87}
88
89fn dft_compl(x: Vec<Complex64>) -> Vec<Complex64> {
90 let mut w = Complex::new(0_f64, -2_f64 * PI / x.len() as f64);
91
92 // f_k (dft_matrix) = SUM{n=0, N-1} f_n * e^(-j2pi*k*n)/N
93 // https://en.wikipedia.org/wiki/Discrete_Fourier_transform
94 let mut dft_matrix: Vec<Vec<Complex64>> = Vec::new();
95 for i in 0..x.len() {
96 let mut f_k: Vec<Complex64> = Vec::new();
97 for j in 0..x.len() {
98 let i_compl = Complex::new(0_f64, i as f64);
99 let j_compl = Complex::new(0_f64, j as f64);
100 let fe = (w * i_compl * j_compl).exp();
101 f_k.push(fe);
102 }
103 dft_matrix.push(f_k.clone());
104 }
105 let r = mul_mv(dft_matrix, x);
106 r
107}
108
109// idft computes the Inverse Discrete Fourier Transform
110pub fn idft(x: &Vec<Complex64>) -> Vec<f64> {
111 let mut w = Complex::new(0_f64, 2_f64 * PI / x.len() as f64);
112
113 // f_k (dft_matrix) = (SUM{n=0, N-1} f_n * e^(j2pi*k*n)/N)/N
114 let mut dft_matrix: Vec<Vec<Complex64>> = Vec::new();
115 for i in 0..x.len() {
116 let mut f_k: Vec<Complex64> = Vec::new();
117 for j in 0..x.len() {
118 let i_compl = Complex::new(0_f64, i as f64);
119 let j_compl = Complex::new(0_f64, j as f64);
120 let fe = (w * i_compl * j_compl).exp();
121 f_k.push(fe);
122 }
123 dft_matrix.push(f_k.clone());
124 }
125 let mut r = mul_mv(dft_matrix, x.clone());
126 let n = x.len() as f64;
127 let mut rr: Vec<f64> = Vec::new();
128 for i in 0..r.len() {
129 r[i] = r[i] / Complex::new(n, 0_f64);
130 rr.push(r[i].re);
131 }
132 rr
133}
134
135// mul_mv multiplies a Matrix by a Vector
136fn mul_mv(a: Vec<Vec<Complex64>>, b: Vec<Complex64>) -> Vec<Complex64> {
137 if a[0].len() != a.len() {
138 panic!("err b[0].len():{:?} != b.len():{:?}", a[0].len(), a.len());
139 }
140 if a.len() != b.len() {
141 panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
142 }
143
144 let rows = a.len();
145 let cols = a.len();
146
147 let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); cols];
148 for i in 0..rows {
149 for j in 0..cols {
150 c[i] += a[i][j] * b[j];
151 }
152 }
153 c
154}
155
156fn add_vv(a: Vec<Complex64>, b: Vec<Complex64>) -> Vec<Complex64> {
157 if a.len() != b.len() {
158 panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
159 }
160 let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); a.len()];
161 for i in 0..a.len() {
162 c[i] = a[i] + b[i];
163 }
164 c
165}
166
167// mul_vv_el multiplies elements of one vector by the elements of another vector
168fn mul_vv_el(a: Vec<Complex64>, b: Vec<Complex64>) -> Vec<Complex64> {
169 if a.len() != b.len() {
170 panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
171 }
172 let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); a.len()];
173 for i in 0..a.len() {
174 c[i] = a[i] * b[i];
175 }
176 c
177}
178
179#[cfg(test)]
180mod tests {
181 use super::*;
182 extern crate rand;
183 use rand::Rng;
184
185 #[test]
186 fn test_dft_simple_values() {
187 let values: Vec<f64> = vec![0.2, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
188 let r = dft(&values);
189 assert_eq!(r.len(), 8);
190
191 assert_eq!(format!("{:.2}", r[0]), "3.70+0.00i");
192 assert_eq!(format!("{:.2}", r[1]), "-0.30-0.97i");
193 assert_eq!(format!("{:.2}", r[2]), "-0.30-0.40i");
194 assert_eq!(format!("{:.2}", r[3]), "-0.30-0.17i");
195 assert_eq!(format!("{:.2}", r[4]), "-0.30+0.00i");
196
197 // expect result similar to initial values
198 let o = idft(&r);
199 assert_eq!(format!("{:.1}", o[0]), "0.2");
200 assert_eq!(format!("{:.1}", o[1]), "0.2");
201 assert_eq!(format!("{:.1}", o[2]), "0.3");
202 assert_eq!(format!("{:.1}", o[3]), "0.4");
203 assert_eq!(format!("{:.1}", o[4]), "0.5");
204 assert_eq!(format!("{:.1}", o[5]), "0.6");
205 assert_eq!(format!("{:.1}", o[6]), "0.7");
206 assert_eq!(format!("{:.1}", o[7]), "0.8");
207 }
208
209 #[test]
210 fn test_fft_simple_values() {
211 let values: Vec<f64> = vec![0.2, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
212 let r = fft(&values);
213
214 assert_eq!(r.len(), 8);
215
216 assert_eq!(format!("{:.2}", r[0]), "3.70+0.00i");
217 assert_eq!(format!("{:.2}", r[1]), "-0.30-0.97i");
218 assert_eq!(format!("{:.2}", r[2]), "-0.30-0.40i");
219 assert_eq!(format!("{:.2}", r[3]), "-0.30-0.17i");
220 assert_eq!(format!("{:.2}", r[4]), "-0.30+0.00i");
221
222 // expect result similar to initial values
223 let o = ifft(&r);
224 println!("{:?}", o);
225 assert_eq!(format!("{:.1}", o[0]), "0.2");
226 assert_eq!(format!("{:.1}", o[1]), "0.2");
227 assert_eq!(format!("{:.1}", o[2]), "0.3");
228 assert_eq!(format!("{:.1}", o[3]), "0.4");
229 assert_eq!(format!("{:.1}", o[4]), "0.5");
230 assert_eq!(format!("{:.1}", o[5]), "0.6");
231 assert_eq!(format!("{:.1}", o[6]), "0.7");
232 assert_eq!(format!("{:.1}", o[7]), "0.8");
233 }
234
235 #[test]
236 fn test_dft_random_values() {
237 let values: Vec<f64> = rand::thread_rng()
238 .sample_iter(rand::distributions::Standard)
239 .take(1024)
240 .collect();
241 let r = dft(&values);
242 println!("{:?}", r.len());
243 let o = idft(&r);
244 }
245
246 #[test]
247 fn test_fft_random_values() {
248 let values: Vec<f64> = rand::thread_rng()
249 .sample_iter(rand::distributions::Standard)
250 .take(1024)
251 .collect();
252 let r = fft(&values);
253 println!("{:?}", r.len());
254 let o = ifft(&r);
255 }
256}