X-Git-Url: https://git.nihav.org/?a=blobdiff_plain;f=nihav-core%2Fsrc%2Fdsp%2Ffft.rs;h=4629f75df5a72dcbdffcc7c07b8d9aa6ada29719;hb=4e034a32d947e1ef5f357cc2477d6f1c3b8454a9;hp=ee4ab7c774b58c065a75b2263126ec40fbc9a7fb;hpb=e243ceb4d694cc08767ad70027bb6963f4cefea3;p=nihav.git diff --git a/nihav-core/src/dsp/fft.rs b/nihav-core/src/dsp/fft.rs index ee4ab7c..4629f75 100644 --- a/nihav-core/src/dsp/fft.rs +++ b/nihav-core/src/dsp/fft.rs @@ -1,21 +1,28 @@ +//! FFT and RDFT implementation. use std::f32::{self, consts}; use std::ops::{Not, Neg, Add, AddAssign, Sub, SubAssign, Mul, MulAssign}; use std::fmt; +/// Complex number. #[repr(C)] #[derive(Debug,Clone,Copy,PartialEq)] pub struct FFTComplex { + /// Real part of the numner. pub re: f32, + /// Complex part of the number. pub im: f32, } impl FFTComplex { + /// Calculates `exp(i * val)`. pub fn exp(val: f32) -> Self { FFTComplex { re: val.cos(), im: val.sin() } } + /// Returns `-Im + i * Re`. pub fn rotate(self) -> Self { FFTComplex { re: -self.im, im: self.re } } + /// Multiplies complex number by scalar. pub fn scale(self, scale: f32) -> Self { FFTComplex { re: self.re * scale, im: self.im * scale } } @@ -86,8 +93,10 @@ impl fmt::Display for FFTComplex { } } +/// Complex number with zero value. pub const FFTC_ZERO: FFTComplex = FFTComplex { re: 0.0, im: 0.0 }; +/// Calculates forward or inverse FFT in the straightforward way. pub fn generic_fft(data: &mut [FFTComplex], forward: bool) { let mut tmp = Vec::with_capacity(data.len()); tmp.resize(data.len(), FFTC_ZERO); @@ -245,7 +254,7 @@ impl FFTSplitRadix { let size = 1 << bits; let mut table = Vec::with_capacity(size); for _ in 0..4 { table.push(FFTC_ZERO); } - for b in 3..(bits+1) { + for b in 3..=bits { let qsize = (1 << (b - 2)) as usize; let base = -consts::PI / ((qsize * 2) as f32); for k in 0..qsize { @@ -523,6 +532,7 @@ impl FFTMode { } } +/// FFT working context. pub struct FFT { perms: Vec, swaps: Vec, @@ -530,14 +540,17 @@ pub struct FFT { } impl FFT { + /// Calculates Fourier transform. pub fn do_fft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) { for k in 0..src.len() { dst[k] = src[self.perms[k]]; } self.do_fft_core(dst); } + /// Calculates inverse Fourier transform. pub fn do_ifft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) { for k in 0..src.len() { dst[k] = src[self.perms[k]]; } self.do_ifft_core(dst); } + /// Performs inplace FFT. pub fn do_fft_inplace(&mut self, data: &mut [FFTComplex]) { for idx in 0..self.swaps.len() { let nidx = self.swaps[idx]; @@ -547,6 +560,7 @@ impl FFT { } self.do_fft_core(data); } + /// Performs inplace inverse FFT. pub fn do_ifft_inplace(&mut self, data: &mut [FFTComplex]) { for idx in 0..self.swaps.len() { let nidx = self.swaps[idx]; @@ -606,6 +620,9 @@ impl FFT { } } +/// [`FFT`] context creator. +/// +/// [`FFT`]: ./struct.FFT.html pub struct FFTBuilder { } @@ -648,7 +665,7 @@ fn gen_sr_perms(swaps: &mut [usize], size: usize) { gen_sr_perms(&mut swaps[3*size/4..], size/4); } -fn gen_swaps_for_perm(swaps: &mut Vec, perms: &Vec) { +fn gen_swaps_for_perm(swaps: &mut Vec, perms: &[usize]) { 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; @@ -683,6 +700,7 @@ impl FFTBuilder { } } } + /// Constructs a new `FFT` context. pub fn new_fft(size: usize, forward: bool) -> FFT { let mut ffts: Vec<(FFTMode, FFTData)> = Vec::with_capacity(1); let mut perms: Vec = Vec::with_capacity(size); @@ -719,12 +737,13 @@ impl FFTBuilder { for (mode, _) in ffts.iter().rev() { mode.permute(&mut perms); } - gen_swaps_for_perm(&mut swaps, &perms); + gen_swaps_for_perm(&mut swaps, perms.as_slice()); FFT { perms, swaps, ffts } } } +/// RDFT working context. pub struct RDFT { table: Vec, fft: FFT, @@ -733,22 +752,24 @@ pub struct RDFT { fwd_fft: bool, } -fn crossadd(a: &FFTComplex, b: &FFTComplex) -> FFTComplex { +fn crossadd(a: FFTComplex, b: FFTComplex) -> FFTComplex { FFTComplex { re: a.re + b.re, im: a.im - b.im } } impl RDFT { + /// Calculates RDFT. pub fn do_rdft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) { dst.copy_from_slice(src); self.do_rdft_inplace(dst); } + /// Calculates inplace RDFT. pub fn do_rdft_inplace(&mut self, buf: &mut [FFTComplex]) { if !self.fwd { for n in 0..self.size/2 { let in0 = buf[n + 1]; let in1 = buf[self.size - n - 1]; - let t0 = crossadd(&in0, &in1); + let t0 = crossadd(in0, in1); let t1 = FFTComplex { re: in1.im + in0.im, im: in1.re - in0.re }; let tab = self.table[n]; let t2 = FFTComplex { re: t1.im * tab.im + t1.re * tab.re, im: t1.im * tab.re - t1.re * tab.im }; @@ -771,12 +792,12 @@ impl RDFT { let in0 = buf[n + 1]; let in1 = buf[self.size - n - 1]; - let t0 = crossadd(&in0, &in1).scale(0.5); + let t0 = crossadd(in0, in1).scale(0.5); let t1 = FFTComplex { re: in0.im + in1.im, im: in0.re - in1.re }; let t2 = t1 * self.table[n]; - buf[n + 1] = crossadd(&t0, &t2); - buf[self.size - n - 1] = FFTComplex { re: t0.re - t2.re, im: -(t0.im + t2.im) }; + buf[n + 1] = crossadd(t0, t2); + buf[self.size - n - 1] = FFTComplex { re: t0.re - t2.re, im: -(t0.im + t2.im) }; } let a = buf[0].re; let b = buf[0].im; @@ -790,10 +811,14 @@ impl RDFT { } } +/// [`RDFT`] context creator. +/// +/// [`RDFT`]: ./struct.FFT.html pub struct RDFTBuilder { } impl RDFTBuilder { + /// Constructs a new `RDFT` context. pub fn new_rdft(size: usize, forward: bool, forward_fft: bool) -> RDFT { let mut table: Vec = Vec::with_capacity(size / 4); let (base, scale) = if forward { (consts::PI / (size as f32), 0.5) } else { (-consts::PI / (size as f32), 1.0) };