dft.c

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

double const PI_2 = 6.28318530718;

float calc_time(unsigned idx, unsigned len) {
  return 2.0 * idx / len - 1;
}

typedef struct Fourier {
  float *real, *imag;
  unsigned max_freq;
} Fourier;

Fourier *make_fourier(unsigned max_freq) {
  Fourier *ptr = (Fourier *)malloc(sizeof(Fourier));
  ptr->real = (float *)calloc(max_freq, sizeof(float));
  ptr->imag = (float *)calloc(max_freq, sizeof(float));
  ptr->max_freq = max_freq;
  return ptr;
}

void drop_fourier(Fourier *f) {
  free(f->imag);
  free(f->real);
  free(f);
}

void fourier_transform(Fourier *f, float const *samples, unsigned samples_len) {
  for (unsigned freq = 0; freq < f->max_freq; ++freq) {
    f->real[freq] = f->imag[freq] = 0;

    float const arg = freq * PI_2 / samples_len;
    for (unsigned time = 0; time < samples_len; ++time) {
      float const sample = samples[time],
        phase = arg * time;
      f->real[freq] += sample * cos(phase);
      f->imag[freq] += sample * sin(phase);
    }
  }
}

void inverse_fourier_transform(Fourier const *f, float *output, unsigned output_len) {
  for (unsigned time = 0; time < output_len; ++time) {
    float real = 0, imag = 0;
    float const arg = time * PI_2 / output_len;
    for (unsigned freq = 0; freq < f->max_freq; ++freq) {
      float const phase = arg * freq;
      real += f->real[freq] * cos(phase) + f->imag[freq] * sin(phase);
      imag -= f->real[freq] * sin(phase) + f->imag[freq] * cos(phase);
    }
    output[time] = sqrt(real * real + imag * imag) / output_len;
  }
}

void display(float const *wave, unsigned len) {
  puts("t,f(t)");
  for (unsigned i = 0; i < len; ++i) {
    printf("%f,%f\n", calc_time(i, len - 1), wave[i]);
  }
}

// wave's domain must be [-1, 1]
void write_with_wave(float (*wave)(float t), float *buf, unsigned buf_len) {
  for (unsigned idx = 0; idx < buf_len; ++idx) {
    buf[idx] = wave(calc_time(idx, buf_len - 1));
  }
}

float saw(float t) {
  return t / 2;
}

float square(float t) {
  return t < 0 ? 0 : 1;
}

int main() {
  float (*const waves[])(float t) = {sinf};
  unsigned const samples_len = 50;
  float input[samples_len] = {}, output[samples_len] = {};
  for (int count = 0; count < 3 * sizeof(waves) / sizeof(waves[0]); ++count) {
    if (count % 3 == 0) {
      write_with_wave(waves[count / 3], input, samples_len);
      puts("original");
      display(input, samples_len);
    }
    int max_freq = 5 * (1 << (count % 3));
    Fourier *f = make_fourier(max_freq);
    fourier_transform(f, input, samples_len);
    inverse_fourier_transform(f, output, samples_len);
    printf("N=%d\n", max_freq);
    display(output, samples_len);
    drop_fourier(f);
  }
  return 0;
}

dft.rs

const PI_2: f64 = 6.28318530718;

fn calc_time(idx: usize, len: usize) -> f64 {
    2.0 * idx as f64 / len as f64 - 1.0
}

struct Fourier {
    real: Vec<f64>,
    imag: Vec<f64>,
    max_freq: usize,
}

impl Fourier {
    fn new(max_freq: usize) -> Self {
        Self {
            real: vec![0.0; max_freq],
            imag: vec![0.0; max_freq],
            max_freq,
        }
    }
    
    fn fourier_transform(&mut self, samples: &[f64]) {
        for freq in 0..self.max_freq {
            self.real[freq] = 0.0;
            self.imag[freq] = 0.0;
            
            let arg = freq as f64 * PI_2 / samples.len() as f64;
            for time in 0..samples.len() {
                let sample = samples[time];
                let phase = arg * time as f64;
                self.real[freq] += sample * phase.cos();
                self.imag[freq] += sample * phase.sin();
            }
        }
    }
    
    fn inv_fourier_transform(&self, output: &mut [f64]) {
        for time in 0..output.len() {
            let mut real = 0.0;
            let mut imag = 0.0;
            let arg = time as f64 * PI_2 / output.len() as f64;
            for freq in 0..self.max_freq {
                let phase = arg * freq as f64;
                real += self.real[freq] * phase.cos();
                real += self.imag[freq] * phase.sin();
                imag -= self.real[freq] * phase.sin();
                imag -= self.imag[freq] * phase.cos();
            }
            output[time] = real.hypot(imag) / output.len() as f64;
        }
    }
}

fn write_with_wave(wave: impl Fn(f64) -> f64, buf: &mut [f64]) {
    for idx in 0..buf.len() {
        buf[idx] = wave(calc_time(idx, buf.len() - 1));
    }
}

fn main() {
    let sin = |t: f64| t.sin();
    let saw = |t| t / 2.0;
    let square = |t| if t < 0.0 { 0.0 } else { 1.0 };
    
    const LEN: usize = 50;
    let mut input = vec![0.0; LEN];
    let mut output = vec![0.0; LEN];
    
    for gen in [sin, saw, square] {
        write_with_wave(gen, &mut input);

        let max_freq = 40;
        let mut f = Fourier::new(max_freq);
        println!("{:?}", input);
        f.fourier_transform(&input);
        f.inv_fourier_transform(&mut output);
        println!("{:?}", output);
    }
}