DFT Trait

src/main.rs

@@ -1,12 +1,21 @@
 use std::{
     f32::consts::PI,
     fs::File,
-    io::Read,
+    io::{Read, Write},
     ops::{Add, Div, Mul, Sub},
 };
 
+mod bfsk;
 mod complex;
+pub mod fft;
+mod nco;
+
+use bfsk::BFSKMod;
 use complex::Complex;
+use complex::Complex32;
+use nco::Nco;
+
+use crate::fft::FFT;
 
 // Utilities
 fn map<T>(input: T, in_min: T, in_max: T, out_min: T, out_max: T) -> T
@@ -20,142 +29,44 @@ fn euclid_mod(a: f32, m: f32) -> f32 {
     let r = a % m;
     if r < 0.0 { r + m } else { r }
 }
-
-struct Nco {
-    // Phase of NCO
-    theta: u32,  // 0 <=> 0, 0xFFFFFFFF <=> 2pi
-    dtheta: u32, // Dtheta = freq : f = dtheta/dt
+fn main() {
+    test();
 }
 
-impl Nco {
-    pub fn new(freq: f32) -> Self {
-        let mut nco = Nco {
-            theta: 0,
-            dtheta: 0,
-        };
-        nco.set_frequency(freq);
-        nco
-    }
+fn test() {
+    let freq1 = 2. * PI / 4.0;
+    let freq2 = 2. * PI / 8.0;
 
-    // Sets freq, freq in radian per sample
-    pub fn set_frequency(&mut self, freq: f32) {
-        self.dtheta = map(euclid_mod(freq, 2. * PI), 0., 2. * PI, 0., 0xFFFFFFFFu32 as f32).floor() as u32;
-    }
-
-    // Adjusts freq, freq in radian per sample
-    pub fn adjust_frequency(&mut self, freq_off: f32) {
-        self.set_frequency(self.get_frequency() + freq_off);
-    }
-
-    pub fn get_frequency(&self) -> f32 {
-        map(self.dtheta as f32, 0., 0xFFFFFFFFu32 as f32, 0., 2. * PI)
-    }
-
-    pub fn step(&mut self) {
-        let bef = self.theta as i64;
-        self.theta = self.theta.overflowing_add(self.dtheta).0;
-    }
-
-    pub fn step_n(&mut self, n: u32) {
-        self.theta = self
-            .theta
-            .overflowing_add(self.dtheta.overflowing_mul(n).0)
-            .0;
-    }
-
-    pub fn sin(&self) -> f32 {
-        map(self.theta as f32, 0., 0xFFFFFFFFu32 as f32, 0., 2. * PI).sin()
-    }
-
-    pub fn cos(&self) -> f32 {
-        map(self.theta as f32, 0., 0xFFFFFFFFu32 as f32, 0., 2. * PI).cos()
-    }
-
-    pub fn cexp(&self) -> Complex<f32>
-    {
-        Complex::new(self.cos(), self.sin())
-    }
-}
-
-struct BFSKMod<'a, T: Iterator<Item = bool>> {
-    samples_per_bit: u32,
-    bandwidth: f32,
-    bit_stream: &'a mut T,
-
-    // State
-    oscillator: Nco,
-    sample_index: u32,
-}
-
-impl<'a, T> BFSKMod<'a, T>
-where
-    T: Iterator<Item = bool>,
-{
-    pub fn new(samples_per_bit: u32, bandwidth: f32, bit_stream: &'a mut T) -> Self {
-        BFSKMod {
-            samples_per_bit,
-            bandwidth,
-            oscillator: Nco::new(0.0),
-            bit_stream,
-            sample_index: samples_per_bit,
-        }
-    }
-
-    pub fn step_modulate(&mut self) -> Option<Complex<f32>> {
-        if self.sample_index == self.samples_per_bit {
-            self.sample_index = 0;
-            let bit = self.bit_stream.next()?;
-
-            let frequency = if bit { self.bandwidth / 2.0 } else { -self.bandwidth / 2.0 };
-            self.oscillator.set_frequency(frequency);
-        }
-
-        self.sample_index += 1;
-        self.oscillator.step();
-        Some(self.oscillator.cexp())
-    }
-}
-
-fn main()
-{
-    modulate();
-}
-
-fn test()
-{
-    let sample_rate = 44100;
-    let f1 = -100.0; //HZ
-    let f2 = 500.0; //HZ
-
-    let spec = hound::WavSpec {
-        channels: 1,
-        sample_rate,
-        bits_per_sample: 16,
-        sample_format: hound::SampleFormat::Int,
-    };
-    let mut writer = hound::WavWriter::create("sine.wav", spec).unwrap();
-
-    let mut o1 = Nco::new(2. * PI * (f1 / sample_rate as f32));
-    let mut o2 = Nco::new(2. * PI * (f2 / sample_rate as f32));
-
-    for i in 0..sample_rate
-    {
-        let amplitude = i16::MAX as f32;
-        let sample = o1.cexp() * o2.cexp();
-        writer.write_sample((amplitude * sample.re) as i16).unwrap();
+    let mut o1 = Nco::new(freq1);
+    let mut o2 = Nco::new(freq2);
 
+    let mut vals = [Complex32::zero(); 8192];
+    for x in vals.iter_mut() {
+        *x = o1.cexp() + o2.cexp();
+        //*x = o2.cexp(); //+ o2.cexp();
+        //*x = *x * (1. / x.mag());
         o1.step();
         o2.step();
     }
 
-    writer.finalize().unwrap();
+    let mut fft = FFT::new(13);
+    let output = fft.run_fft(&vals);
+
+    let mut f = File::create("out.csv").unwrap();
+    for (i, v) in output.iter().enumerate() {
+        f.write_all(
+            format!("{},{},{},\n", i as f32 / 8192., v.mag(), v.arg())
+                .to_string()
+                .as_bytes(),
+        )
+        .unwrap();
+    }
 }
 
 fn modulate() {
     let sample_rate = 44100;
     let mut frequency = 2000.0; //HZ
     let mut bandwidth = 500.0; //HZ
-                               
 
     let path = "a.jpg";
     let file = File::open(path).unwrap();
@@ -176,7 +87,7 @@ fn modulate() {
     //let mut bit_stream = (0..22000).flat_map(|_| [true, false]);
 
     let baud_rate = 400;
-    println!("{} samples/bit", sample_rate/baud_rate);
+    println!("{} samples/bit", sample_rate / baud_rate);
     let mut bfsk = BFSKMod::new(
         sample_rate / baud_rate,
         2. * PI * (bandwidth / sample_rate as f32),
@@ -200,7 +111,9 @@ fn modulate() {
         let c_sample = lo.cexp() * sample;
 
         let filtered = prev + (c_sample - prev) * alpha;
-        writer.write_sample((amplitude * c_sample.re) as i16).unwrap();
+        writer
+            .write_sample((amplitude * c_sample.re) as i16)
+            .unwrap();
         lo.step();
     }
     writer.finalize().unwrap();