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_mv(f, x_compl);
 27    r
 28}
 29
 30// idft computes the Inverse Discrete Fourier Transform
 31pub fn idft(x: &Vec<Complex64>) -> Vec<f64> {
 32    let mut w = Complex::new(0_f64, 2_f64 * PI / x.len() as f64);
 33
 34    // f_k = (SUM{n=0, N-1} f_n * e^(j2pi*k*n)/N)/N
 35    let mut dft_matrix: Vec<Vec<Complex64>> = Vec::new();
 36    for i in 0..x.len() {
 37        let mut f_k: Vec<Complex64> = Vec::new();
 38        for j in 0..x.len() {
 39            let i_compl = Complex::new(0_f64, i as f64);
 40            let j_compl = Complex::new(0_f64, j as f64);
 41            let fe = (w * i_compl * j_compl).exp();
 42            f_k.push(fe);
 43        }
 44        dft_matrix.push(f_k.clone());
 45    }
 46    let mut r = mul_mv(dft_matrix, x.clone());
 47    let n = x.len() as f64;
 48    let mut rr: Vec<f64> = Vec::new();
 49    for i in 0..r.len() {
 50        r[i] = r[i] / Complex::new(n, 0_f64);
 51        rr.push(r[i].re);
 52    }
 53    rr
 54}
 55
 56// mul_mv multiplies a Matrix by a Vector
 57fn mul_mv(a: Vec<Vec<Complex64>>, b: Vec<Complex64>) -> Vec<Complex64> {
 58    if a[0].len() != a.len() {
 59        panic!("err b[0].len():{:?} != b.len():{:?}", a[0].len(), a.len());
 60    }
 61    if a.len() != b.len() {
 62        panic!("err a.len():{:?} != b.len():{:?}", a.len(), b.len());
 63    }
 64
 65    let rows = a.len();
 66    let cols = a.len();
 67
 68    let mut c: Vec<Complex64> = vec![Complex::new(0_f64, 0_f64); cols];
 69    for i in 0..rows {
 70        for j in 0..cols {
 71            c[i] += a[i][j] * b[j];
 72        }
 73    }
 74    c
 75}
 76
 77#[cfg(test)]
 78mod tests {
 79    use super::*;
 80
 81    #[test]
 82    fn test_dft_simple_values() {
 83        let values: Vec<f64> = vec![0.2, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8];
 84        let r = dft(&values);
 85        assert_eq!(r.len(), 8);
 86
 87        assert_eq!(format!("{:.2}", r[0]), "3.70+0.00i");
 88        assert_eq!(format!("{:.2}", r[1]), "-0.30-0.97i");
 89        assert_eq!(format!("{:.2}", r[2]), "-0.30-0.40i");
 90        assert_eq!(format!("{:.2}", r[3]), "-0.30-0.17i");
 91        assert_eq!(format!("{:.2}", r[4]), "-0.30+0.00i");
 92
 93        // expect result similar to initial values
 94        let o = idft(&r);
 95        assert_eq!(format!("{:.1}", o[0]), "0.2");
 96        assert_eq!(format!("{:.1}", o[1]), "0.2");
 97        assert_eq!(format!("{:.1}", o[2]), "0.3");
 98        assert_eq!(format!("{:.1}", o[3]), "0.4");
 99        assert_eq!(format!("{:.1}", o[4]), "0.5");
100        assert_eq!(format!("{:.1}", o[5]), "0.6");
101        assert_eq!(format!("{:.1}", o[6]), "0.7");
102        assert_eq!(format!("{:.1}", o[7]), "0.8");
103    }
104}