1use num::complex::{Complex, Complex64};
 2use std::f64::consts::PI;
 3
 4// dft computes the Discrete Fourier Transform
 5pub fn dft(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
11    let mut w = Complex::new(0_f64, -2_f64 * PI / x.len() as f64);
12
13    // f_k = SUM{n=0, N-1} f_n * e^(-j2pi*k*n)/N
14    // https://en.wikipedia.org/wiki/Discrete_Fourier_transform
15    let mut f: Vec<Vec<Complex64>> = Vec::new();
16    for i in 0..x.len() {
17        let mut f_k: Vec<Complex64> = Vec::new();
18        for j in 0..x.len() {
19            let i_compl = Complex::new(0_f64, i as f64);
20            let j_compl = Complex::new(0_f64, j as f64);
21            let fe = (w * i_compl * j_compl).exp();
22            f_k.push(fe);
23        }
24        f.push(f_k.clone());
25    }
26    let r = mul_vv(f, x_compl);
27    r
28}
29
30// mul_vv multiplies a Matrix by a Vector
31fn mul_vv(a: Vec<Vec<Complex64>>, b: Vec<Complex64>) -> Vec<Complex64> {
32    if a[0].len() != a.len() {
33        panic!("err b[0].len():{:?} != b.len():{:?}", a[0].len(), a.len());
34    }
35    if a.len() != b.len() {
36        panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
37    }
38
39    let rows = a.len();
40    let cols = a.len();
41
42    let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); cols];
43    for i in 0..rows {
44        for j in 0..cols {
45            c[i] += a[i][j] * b[j];
46        }
47    }
48    c
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn test_dft_simple_values() {
57        let values: Vec<f64> = vec![0.2, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
58        let r = dft(&values);
59        assert_eq!(r.len(), 8);
60
61        assert_eq!(format!("{:.2}", r[0]), "3.70+0.00i");
62        assert_eq!(format!("{:.2}", r[1]), "-0.30-0.97i");
63        assert_eq!(format!("{:.2}", r[2]), "-0.30-0.40i");
64        assert_eq!(format!("{:.2}", r[3]), "-0.30-0.17i");
65        assert_eq!(format!("{:.2}", r[4]), "-0.30+0.00i");
66    }
67}