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 N = x.len();
7 if N % 2 > 0 {
8 panic!("not a power of 2");
9 } else if N <= 2 {
10 return dft(x);
11 }
12
13 let mut x_even: Vec<f64> = Vec::new();
14 let mut x_odd: Vec<f64> = Vec::new();
15 for i in 0..x.len() {
16 if i % 2 == 0 {
17 x_even.push(x[i]);
18 } else {
19 x_odd.push(x[i]);
20 }
21 }
22 let mut x_even_cmplx = fft(&x_even);
23 let mut x_odd_cmplx = fft(&x_odd);
24
25 let mut w = Complex::new(0_f64, 2_f64 * PI / N as f64);
26 let mut f_k: Vec<Complex64> = Vec::new();
27 for k in 0..x.len() {
28 let k_compl = Complex::new(k as f64, 0_f64);
29 f_k.push((w * k_compl).exp());
30 }
31
32 let mut r: Vec<Complex64> = Vec::new();
33 let mut aa = add_vv(
34 x_even_cmplx.clone(),
35 mul_vv_el(x_odd_cmplx.clone(), f_k.clone()[0..x.len() / 2].to_vec()),
36 );
37 let mut bb = add_vv(
38 x_even_cmplx.clone(),
39 mul_vv_el(x_odd_cmplx.clone(), f_k.clone()[x.len() / 2..].to_vec()),
40 );
41 r.append(&mut aa);
42 r.append(&mut bb);
43
44 r
45}
46
47// dft computes the Discrete Fourier Transform
48pub fn dft(x: &Vec<f64>) -> Vec<Complex64> {
49 let mut x_compl: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); x.len()];
50 for i in 0..x.len() {
51 x_compl[i] = Complex::new(x[i], 0_f64);
52 }
53
54 let mut w = Complex::new(0_f64, -2_f64 * PI / x.len() as f64);
55
56 // f_k = SUM{n=0, N-1} f_n * e^(-j2pi*k*n)/N
57 // https://en.wikipedia.org/wiki/Discrete_Fourier_transform
58 let mut f: Vec<Vec<Complex64>> = Vec::new();
59 for i in 0..x.len() {
60 let mut f_k: Vec<Complex64> = Vec::new();
61 for j in 0..x.len() {
62 let i_compl = Complex::new(0_f64, i as f64);
63 let j_compl = Complex::new(0_f64, j as f64);
64 let fe = (w * i_compl * j_compl).exp();
65 f_k.push(fe);
66 }
67 f.push(f_k.clone());
68 }
69 let r = mul_mv(f, x_compl);
70 r
71}
72
73// idft computes the Inverse Discrete Fourier Transform
74pub fn idft(x: &Vec<Complex64>) -> Vec<f64> {
75 let mut w = Complex::new(0_f64, 2_f64 * PI / x.len() as f64);
76
77 // f_k = (SUM{n=0, N-1} f_n * e^(j2pi*k*n)/N)/N
78 let mut dft_matrix: Vec<Vec<Complex64>> = Vec::new();
79 for i in 0..x.len() {
80 let mut f_k: Vec<Complex64> = Vec::new();
81 for j in 0..x.len() {
82 let i_compl = Complex::new(0_f64, i as f64);
83 let j_compl = Complex::new(0_f64, j as f64);
84 let fe = (w * i_compl * j_compl).exp();
85 f_k.push(fe);
86 }
87 dft_matrix.push(f_k.clone());
88 }
89 let mut r = mul_mv(dft_matrix, x.clone());
90 let n = x.len() as f64;
91 let mut rr: Vec<f64> = Vec::new();
92 for i in 0..r.len() {
93 r[i] = r[i] / Complex::new(n, 0_f64);
94 rr.push(r[i].re);
95 }
96 rr
97}
98
99// mul_mv multiplies a Matrix by a Vector
100fn mul_mv(a: Vec<Vec<Complex64>>, b: Vec<Complex64>) -> Vec<Complex64> {
101 if a[0].len() != a.len() {
102 panic!("err b[0].len():{:?} != b.len():{:?}", a[0].len(), a.len());
103 }
104 if a.len() != b.len() {
105 panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
106 }
107
108 let rows = a.len();
109 let cols = a.len();
110
111 let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); cols];
112 for i in 0..rows {
113 for j in 0..cols {
114 c[i] += a[i][j] * b[j];
115 }
116 }
117 c
118}
119
120fn add_vv(a: Vec<Complex64>, b: Vec<Complex64>) -> Vec<Complex64> {
121 if a.len() != b.len() {
122 panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
123 }
124 let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); a.len()];
125 for i in 0..a.len() {
126 c[i] = a[i] + b[i];
127 }
128 c
129}
130
131// mul_vv_el multiplies elements of one vector by the elements of another vector
132fn mul_vv_el(a: Vec<Complex64>, b: Vec<Complex64>) -> Vec<Complex64> {
133 if a.len() != b.len() {
134 panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
135 }
136 let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); a.len()];
137 for i in 0..a.len() {
138 c[i] = a[i] * b[i];
139 }
140 c
141}
142
143#[cfg(test)]
144mod tests {
145 use super::*;
146
147 #[test]
148 fn test_dft_simple_values() {
149 let values: Vec<f64> = vec![0.2, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
150 let r = dft(&values);
151 assert_eq!(r.len(), 8);
152
153 assert_eq!(format!("{:.2}", r[0]), "3.70+0.00i");
154 assert_eq!(format!("{:.2}", r[1]), "-0.30-0.97i");
155 assert_eq!(format!("{:.2}", r[2]), "-0.30-0.40i");
156 assert_eq!(format!("{:.2}", r[3]), "-0.30-0.17i");
157 assert_eq!(format!("{:.2}", r[4]), "-0.30+0.00i");
158
159 // expect result similar to initial values
160 let o = idft(&r);
161 assert_eq!(format!("{:.1}", o[0]), "0.2");
162 assert_eq!(format!("{:.1}", o[1]), "0.2");
163 assert_eq!(format!("{:.1}", o[2]), "0.3");
164 assert_eq!(format!("{:.1}", o[3]), "0.4");
165 assert_eq!(format!("{:.1}", o[4]), "0.5");
166 assert_eq!(format!("{:.1}", o[5]), "0.6");
167 assert_eq!(format!("{:.1}", o[6]), "0.7");
168 assert_eq!(format!("{:.1}", o[7]), "0.8");
169 }
170
171 #[test]
172 fn test_fft_simple_values() {
173 let values: Vec<f64> = vec![0.2, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
174 let r = fft(&values);
175
176 assert_eq!(r.len(), 8);
177
178 assert_eq!(format!("{:.2}", r[0]), "3.70+0.00i");
179 assert_eq!(format!("{:.2}", r[1]), "-0.30-0.97i");
180 assert_eq!(format!("{:.2}", r[2]), "-0.30-0.40i");
181 assert_eq!(format!("{:.2}", r[3]), "-0.30-0.17i");
182 assert_eq!(format!("{:.2}", r[4]), "-0.30+0.00i");
183 }
184}