dft.rs

use std::ops::{Add,Sub,Mul,Div,Neg};
use std::fmt;
use std::fmt::Display;

#[derive(Clone, Copy, Debug)]
struct Complex {
    re: f32,
    im: f32,
}

impl Complex {
    fn magn(self) -> f32 {
        self.re * self.re + self.im * self.im
    }
}

impl Neg for Complex {
    type Output = Complex;

    fn neg(self) -> Complex {
        Complex {
            re: -self.re,
            im: -self.im,
        }
    }
}

impl Add for Complex {
    type Output = Complex;

    fn add(self, rhs: Complex) -> Complex {
        Complex {
            re: self.re + rhs.re,
            im: self.im + rhs.im,
        }
    }
}

impl Sub for Complex {
    type Output = Complex;

    fn sub(self, rhs: Complex) -> Complex {
        Complex {
            re: self.re - rhs.re,
            im: self.im - rhs.im,
        }
    }
}

impl Mul for Complex {
    type Output = Complex;

    fn mul(self, rhs: Complex) -> Complex {
        Complex {
            re: self.re * rhs.re - self.im * rhs.im,
            im: self.re * rhs.im + self.im * rhs.re,
        }
    }
}

impl Div for Complex {
    type Output = Complex;

    fn div(self, rhs: Complex) -> Complex {
        let magn = rhs.magn();

        Complex {
            re: (self.re * rhs.re + self.im * rhs.im) / magn,
            im: (self.im * rhs.re - self.re * rhs.im) / magn,
        }
    }
}

impl Display for Complex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f, "{:6.3} {} {:.3}i",
            self.re,
            if self.im < 0.0 { "-" } else { "+" },
            self.im.abs()
        )
    }
}

fn exp(z: Complex) -> Complex {
    let magn = z.re.exp();

    Complex {
        re: magn * z.im.cos(),
        im: magn * z.im.sin(),
    }
}

fn dft(x: &Vec<Vec<f32>>, n: usize, k: usize) -> Complex {
    let N = x.len();
    let K = x[0].len();

    assert!(n < N);
    assert!(k < K);

    let mut s = Complex { re: 0.0, im: 0.0 };
    let tauim = Complex { re: 0.0, im: -2.0 * std::f32::consts::PI };

    for l in 0..N {
        for m in 0..K {
            let re = l as f32 * n as f32 / N as f32 + m as f32 * k as f32 / K as f32;
            s = s + Complex { re: x[l][m], im: 0.0 } * exp(tauim * Complex { re: re, im: 0.0 })
        }
    }

    s // Complex { re: N as f32 * K as f32, im: 0.0 }
}

fn main() {
    let v = vec![
        vec![ 1.0,  -1.0,   1.0],
        vec![-1.0,   1.0,  -1.0],
        vec![ 1.0,  -1.0,   1.0],
    ];

    for n in 0..v.len() {
        for k in 0..v[n].len() {
            print!("      {}", dft(&v, n, k))
        }
        println!("");
    }
}