X-Git-Url: https://git.nihav.org/?a=blobdiff_plain;f=src%2Fdsp%2Ffft.rs;h=a626ad95a638b99b2b3a055a72e4d7ab31c8a7cb;hb=HEAD;hp=7f9b4aa75d645500526e75aef604bd28491c76da;hpb=3f3b5b2356c82357a0e29e37cb94b6e77c2a519a;p=nihav.git diff --git a/src/dsp/fft.rs b/src/dsp/fft.rs deleted file mode 100644 index 7f9b4aa..0000000 --- a/src/dsp/fft.rs +++ /dev/null @@ -1,475 +0,0 @@ -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, - perms: Vec, - swaps: Vec, - 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 = 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 = Vec::with_capacity(size / 2); - let mut ovec1: Vec = Vec::with_capacity(size / 4); - let mut ovec2: Vec = 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, perms: &Vec) { - let mut idx_arr: Vec = 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; - let mut perms: Vec; - let mut table: Vec; - 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); - } - } -}