+++ /dev/null
-use std::f32::{self, consts};
-use std::ops::{Not, Neg, Add, AddAssign, Sub, SubAssign, Mul, MulAssign};
-use std::fmt;
-
-#[derive(Debug,Clone,Copy,PartialEq)]
-pub struct FFTComplex {
- pub re: f32,
- pub im: f32,
-}
-
-impl FFTComplex {
- pub fn exp(val: f32) -> Self {
- FFTComplex { re: val.cos(), im: val.sin() }
- }
- pub fn rotate(self) -> Self {
- FFTComplex { re: -self.im, im: self.re }
- }
- pub fn scale(self, scale: f32) -> Self {
- FFTComplex { re: self.re * scale, im: self.im * scale }
- }
-}
-
-impl Neg for FFTComplex {
- type Output = FFTComplex;
- fn neg(self) -> Self::Output {
- FFTComplex { re: -self.re, im: -self.im }
- }
-}
-
-impl Not for FFTComplex {
- type Output = FFTComplex;
- fn not(self) -> Self::Output {
- FFTComplex { re: self.re, im: -self.im }
- }
-}
-
-impl Add for FFTComplex {
- type Output = FFTComplex;
- fn add(self, other: Self) -> Self::Output {
- FFTComplex { re: self.re + other.re, im: self.im + other.im }
- }
-}
-
-impl AddAssign for FFTComplex {
- fn add_assign(&mut self, other: Self) {
- self.re += other.re;
- self.im += other.im;
- }
-}
-
-impl Sub for FFTComplex {
- type Output = FFTComplex;
- fn sub(self, other: Self) -> Self::Output {
- FFTComplex { re: self.re - other.re, im: self.im - other.im }
- }
-}
-
-impl SubAssign for FFTComplex {
- fn sub_assign(&mut self, other: Self) {
- self.re -= other.re;
- self.im -= other.im;
- }
-}
-
-impl Mul for FFTComplex {
- type Output = FFTComplex;
- fn mul(self, other: Self) -> Self::Output {
- FFTComplex { re: self.re * other.re - self.im * other.im,
- im: self.im * other.re + self.re * other.im }
- }
-}
-
-impl MulAssign for FFTComplex {
- fn mul_assign(&mut self, other: Self) {
- let re = self.re * other.re - self.im * other.im;
- let im = self.im * other.re + self.re * other.im;
- self.re = re;
- self.im = im;
- }
-}
-
-impl fmt::Display for FFTComplex {
- fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
- write!(f, "({}, {})", self.re, self.im)
- }
-}
-
-pub const FFTC_ZERO: FFTComplex = FFTComplex { re: 0.0, im: 0.0 };
-
-#[derive(Debug,Clone,Copy,PartialEq)]
-pub enum FFTMode {
- Matrix,
- CooleyTukey,
- SplitRadix,
-}
-
-pub struct FFT {
- table: Vec<FFTComplex>,
- perms: Vec<usize>,
- swaps: Vec<usize>,
- bits: u32,
- mode: FFTMode,
-}
-
-impl FFT {
- fn do_fft_inplace_ct(&mut self, data: &mut [FFTComplex], bits: u32, forward: bool) {
- if bits == 0 { return; }
- if bits == 1 {
- let sum01 = data[0] + data[1];
- let dif01 = data[0] - data[1];
- data[0] = sum01;
- data[1] = dif01;
- return;
- }
- if bits == 2 {
- let sum01 = data[0] + data[1];
- let dif01 = data[0] - data[1];
- let sum23 = data[2] + data[3];
- let dif23 = data[2] - data[3];
- if forward {
- data[0] = sum01 + sum23;
- data[1] = dif01 - dif23.rotate();
- data[2] = sum01 - sum23;
- data[3] = dif01 + dif23.rotate();
- } else {
- data[0] = sum01 + sum23;
- data[1] = dif01 + dif23.rotate();
- data[2] = sum01 - sum23;
- data[3] = dif01 - dif23.rotate();
- }
- return;
- }
-
- let hsize = (1 << (bits - 1)) as usize;
- self.do_fft_inplace_ct(&mut data[0..hsize], bits - 1, forward);
- self.do_fft_inplace_ct(&mut data[hsize..], bits - 1, forward);
- let offs = hsize;
- {
- let e = data[0];
- let o = data[hsize];
- data[0] = e + o;
- data[hsize] = e - o;
- }
- if forward {
- for k in 1..hsize {
- let e = data[k];
- let o = data[k + hsize] * self.table[offs + k];
- data[k] = e + o;
- data[k + hsize] = e - o;
- }
- } else {
- for k in 1..hsize {
- let e = data[k];
- let o = data[k + hsize] * !self.table[offs + k];
- data[k] = e + o;
- data[k + hsize] = e - o;
- }
- }
- }
-
- fn do_fft_inplace_splitradix(&mut self, data: &mut [FFTComplex], bits: u32, forward: bool) {
- if bits == 0 { return; }
- if bits == 1 {
- let sum01 = data[0] + data[1];
- let dif01 = data[0] - data[1];
- data[0] = sum01;
- data[1] = dif01;
- return;
- }
- if bits == 2 {
- let sum01 = data[0] + data[2];
- let dif01 = data[0] - data[2];
- let sum23 = data[1] + data[3];
- let dif23 = data[1] - data[3];
- if forward {
- data[0] = sum01 + sum23;
- data[1] = dif01 - dif23.rotate();
- data[2] = sum01 - sum23;
- data[3] = dif01 + dif23.rotate();
- } else {
- data[0] = sum01 + sum23;
- data[1] = dif01 + dif23.rotate();
- data[2] = sum01 - sum23;
- data[3] = dif01 - dif23.rotate();
- }
- return;
- }
- let qsize = (1 << (bits - 2)) as usize;
- let hsize = (1 << (bits - 1)) as usize;
- let q3size = qsize + hsize;
-
- self.do_fft_inplace_splitradix(&mut data[0 ..hsize], bits - 1, forward);
- self.do_fft_inplace_splitradix(&mut data[hsize ..q3size], bits - 2, forward);
- self.do_fft_inplace_splitradix(&mut data[q3size..], bits - 2, forward);
- let off = hsize;
- if forward {
- {
- let t3 = data[0 + hsize] + data[0 + q3size];
- let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
- let e1 = data[0];
- let e2 = data[0 + qsize];
- data[0] = e1 + t3;
- data[0 + qsize] = e2 - t4;
- data[0 + hsize] = e1 - t3;
- data[0 + q3size] = e2 + t4;
- }
- for k in 1..qsize {
- let t1 = self.table[off + k * 2 + 0] * data[k + hsize];
- let t2 = self.table[off + k * 2 + 1] * data[k + q3size];
- let t3 = t1 + t2;
- let t4 = (t1 - t2).rotate();
- let e1 = data[k];
- let e2 = data[k + qsize];
- data[k] = e1 + t3;
- data[k + qsize] = e2 - t4;
- data[k + hsize] = e1 - t3;
- data[k + qsize * 3] = e2 + t4;
- }
- } else {
- {
- let t3 = data[0 + hsize] + data[0 + q3size];
- let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
- let e1 = data[0];
- let e2 = data[0 + qsize];
- data[0] = e1 + t3;
- data[0 + qsize] = e2 + t4;
- data[0 + hsize] = e1 - t3;
- data[0 + q3size] = e2 - t4;
- }
- for k in 1..qsize {
- let t1 = !self.table[off + k * 2 + 0] * data[k + hsize];
- let t2 = !self.table[off + k * 2 + 1] * data[k + q3size];
- let t3 = t1 + t2;
- let t4 = (t1 - t2).rotate();
- let e1 = data[k];
- let e2 = data[k + qsize];
- data[k] = e1 + t3;
- data[k + qsize] = e2 + t4;
- data[k + hsize] = e1 - t3;
- data[k + qsize * 3] = e2 - t4;
- }
- }
- }
-
- pub fn do_fft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex], forward: bool) {
- match self.mode {
- FFTMode::Matrix => {
- let base = if forward { -consts::PI * 2.0 / (src.len() as f32) }
- else { consts::PI * 2.0 / (src.len() as f32) };
- for k in 0..src.len() {
- let mut sum = FFTC_ZERO;
- for n in 0..src.len() {
- let w = FFTComplex::exp(base * ((n * k) as f32));
- sum += src[n] * w;
- }
- dst[k] = sum;
- }
- },
- FFTMode::CooleyTukey => {
- let bits = self.bits;
- for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
- self.do_fft_inplace_ct(dst, bits, forward);
- },
- FFTMode::SplitRadix => {
- let bits = self.bits;
- for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
- self.do_fft_inplace_splitradix(dst, bits, forward);
- },
- };
- }
-
- pub fn do_fft_inplace(&mut self, data: &mut [FFTComplex], forward: bool) {
- for idx in 0..self.swaps.len() {
- let nidx = self.swaps[idx];
- if idx != nidx {
- let t = data[nidx];
- data[nidx] = data[idx];
- data[idx] = t;
- }
- }
- match self.mode {
- FFTMode::Matrix => {
- let size = (1 << self.bits) as usize;
- let base = if forward { -consts::PI * 2.0 / (size as f32) }
- else { consts::PI * 2.0 / (size as f32) };
- let mut res: Vec<FFTComplex> = Vec::with_capacity(size);
- for k in 0..size {
- let mut sum = FFTC_ZERO;
- for n in 0..size {
- let w = FFTComplex::exp(base * ((n * k) as f32));
- sum += data[n] * w;
- }
- res.push(sum);
- }
- for k in 0..size {
- data[k] = res[k];
- }
- },
- FFTMode::CooleyTukey => {
- let bits = self.bits;
- self.do_fft_inplace_ct(data, bits, forward);
- },
- FFTMode::SplitRadix => {
- let bits = self.bits;
- self.do_fft_inplace_splitradix(data, bits, forward);
- },
- };
- }
-}
-
-pub struct FFTBuilder {
-}
-
-fn reverse_bits(inval: u32) -> u32 {
- const REV_TAB: [u8; 16] = [
- 0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
- 0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
- ];
-
- let mut ret = 0;
- let mut val = inval;
- for _ in 0..8 {
- ret = (ret << 4) | (REV_TAB[(val & 0xF) as usize] as u32);
- val = val >> 4;
- }
- ret
-}
-
-fn swp_idx(idx: usize, bits: u32) -> usize {
- let s = reverse_bits(idx as u32) as usize;
- s >> (32 - bits)
-}
-
-fn gen_sr_perms(swaps: &mut [usize], size: usize) {
- if size <= 4 { return; }
- let mut evec: Vec<usize> = Vec::with_capacity(size / 2);
- let mut ovec1: Vec<usize> = Vec::with_capacity(size / 4);
- let mut ovec2: Vec<usize> = Vec::with_capacity(size / 4);
- for k in 0..size/4 {
- evec.push (swaps[k * 4 + 0]);
- ovec1.push(swaps[k * 4 + 1]);
- evec.push (swaps[k * 4 + 2]);
- ovec2.push(swaps[k * 4 + 3]);
- }
- for k in 0..size/2 { swaps[k] = evec[k]; }
- for k in 0..size/4 { swaps[k + size/2] = ovec1[k]; }
- for k in 0..size/4 { swaps[k + 3*size/4] = ovec2[k]; }
- gen_sr_perms(&mut swaps[0..size/2], size/2);
- gen_sr_perms(&mut swaps[size/2..3*size/4], size/4);
- gen_sr_perms(&mut swaps[3*size/4..], size/4);
-}
-
-fn gen_swaps_for_perm(swaps: &mut Vec<usize>, perms: &Vec<usize>) {
- let mut idx_arr: Vec<usize> = Vec::with_capacity(perms.len());
- for i in 0..perms.len() { idx_arr.push(i); }
- let mut run_size = 0;
- let mut run_pos = 0;
- for idx in 0..perms.len() {
- if perms[idx] == idx_arr[idx] {
- if run_size == 0 { run_pos = idx; }
- run_size += 1;
- } else {
- for i in 0..run_size {
- swaps.push(run_pos + i);
- }
- run_size = 0;
- let mut spos = idx + 1;
- while idx_arr[spos] != perms[idx] { spos += 1; }
- idx_arr[spos] = idx_arr[idx];
- idx_arr[idx] = perms[idx];
- swaps.push(spos);
- }
- }
-}
-
-impl FFTBuilder {
- pub fn new_fft(mode: FFTMode, size: usize) -> FFT {
- let mut swaps: Vec<usize>;
- let mut perms: Vec<usize>;
- let mut table: Vec<FFTComplex>;
- let bits = 31 - (size as u32).leading_zeros();
- match mode {
- FFTMode::Matrix => {
- swaps = Vec::new();
- perms = Vec::new();
- table = Vec::new();
- },
- FFTMode::CooleyTukey => {
- perms = Vec::with_capacity(size);
- for i in 0..size {
- perms.push(swp_idx(i, bits));
- }
- swaps = Vec::with_capacity(size);
- table = Vec::with_capacity(size);
- for _ in 0..4 { table.push(FFTC_ZERO); }
- for b in 3..(bits+1) {
- let hsize = (1 << (b - 1)) as usize;
- let base = -consts::PI / (hsize as f32);
- for k in 0..hsize {
- table.push(FFTComplex::exp(base * (k as f32)));
- }
- }
- },
- FFTMode::SplitRadix => {
- perms = Vec::with_capacity(size);
- for i in 0..size {
- perms.push(i);
- }
- gen_sr_perms(perms.as_mut_slice(), 1 << bits);
- swaps = Vec::with_capacity(size);
- table = Vec::with_capacity(size);
- for _ in 0..4 { table.push(FFTC_ZERO); }
- for b in 3..(bits+1) {
- let qsize = (1 << (b - 2)) as usize;
- let base = -consts::PI / ((qsize * 2) as f32);
- for k in 0..qsize {
- table.push(FFTComplex::exp(base * ((k * 1) as f32)));
- table.push(FFTComplex::exp(base * ((k * 3) as f32)));
- }
- }
- },
- };
- gen_swaps_for_perm(&mut swaps, &perms);
- FFT { mode: mode, swaps: swaps, perms: perms, bits: bits, table: table }
- }
-}
-
-
-#[cfg(test)]
-mod test {
- use super::*;
-
- #[test]
- fn test_fft() {
- let mut fin: [FFTComplex; 128] = [FFTC_ZERO; 128];
- let mut fout1: [FFTComplex; 128] = [FFTC_ZERO; 128];
- let mut fout2: [FFTComplex; 128] = [FFTC_ZERO; 128];
- let mut fout3: [FFTComplex; 128] = [FFTC_ZERO; 128];
- let mut fft1 = FFTBuilder::new_fft(FFTMode::Matrix, fin.len());
- let mut fft2 = FFTBuilder::new_fft(FFTMode::CooleyTukey, fin.len());
- let mut fft3 = FFTBuilder::new_fft(FFTMode::SplitRadix, fin.len());
- let mut seed: u32 = 42;
- for i in 0..fin.len() {
- seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
- let val = (seed >> 16) as i16;
- fin[i].re = (val as f32) / 256.0;
- seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
- let val = (seed >> 16) as i16;
- fin[i].im = (val as f32) / 256.0;
- }
- fft1.do_fft(&fin, &mut fout1, true);
- fft2.do_fft(&fin, &mut fout2, true);
- fft3.do_fft(&fin, &mut fout3, true);
-
- for i in 0..fin.len() {
- assert!((fout1[i].re - fout2[i].re).abs() < 1.0);
- assert!((fout1[i].im - fout2[i].im).abs() < 1.0);
- assert!((fout1[i].re - fout3[i].re).abs() < 1.0);
- assert!((fout1[i].im - fout3[i].im).abs() < 1.0);
- }
- fft1.do_fft_inplace(&mut fout1, false);
- fft2.do_fft_inplace(&mut fout2, false);
- fft3.do_fft_inplace(&mut fout3, false);
-
- let sc = 1.0 / (fin.len() as f32);
- for i in 0..fin.len() {
- assert!((fin[i].re - fout1[i].re * sc).abs() < 1.0);
- assert!((fin[i].im - fout1[i].im * sc).abs() < 1.0);
- assert!((fout1[i].re - fout2[i].re).abs() < 1.0);
- assert!((fout1[i].im - fout2[i].im).abs() < 1.0);
- assert!((fout1[i].re - fout3[i].re).abs() < 1.0);
- assert!((fout1[i].im - fout3[i].im).abs() < 1.0);
- }
- }
-}
projects / nihav.git / blobdiff