1 use std::f32::{self, consts};
2 use std::ops::{Not, Neg, Add, AddAssign, Sub, SubAssign, Mul, MulAssign};
6 #[derive(Debug,Clone,Copy,PartialEq)]
7 pub struct FFTComplex {
13 pub fn exp(val: f32) -> Self {
14 FFTComplex { re: val.cos(), im: val.sin() }
16 pub fn rotate(self) -> Self {
17 FFTComplex { re: -self.im, im: self.re }
19 pub fn scale(self, scale: f32) -> Self {
20 FFTComplex { re: self.re * scale, im: self.im * scale }
24 impl Neg for FFTComplex {
25 type Output = FFTComplex;
26 fn neg(self) -> Self::Output {
27 FFTComplex { re: -self.re, im: -self.im }
31 impl Not for FFTComplex {
32 type Output = FFTComplex;
33 fn not(self) -> Self::Output {
34 FFTComplex { re: self.re, im: -self.im }
38 impl Add for FFTComplex {
39 type Output = FFTComplex;
40 fn add(self, other: Self) -> Self::Output {
41 FFTComplex { re: self.re + other.re, im: self.im + other.im }
45 impl AddAssign for FFTComplex {
46 fn add_assign(&mut self, other: Self) {
52 impl Sub for FFTComplex {
53 type Output = FFTComplex;
54 fn sub(self, other: Self) -> Self::Output {
55 FFTComplex { re: self.re - other.re, im: self.im - other.im }
59 impl SubAssign for FFTComplex {
60 fn sub_assign(&mut self, other: Self) {
66 impl Mul for FFTComplex {
67 type Output = FFTComplex;
68 fn mul(self, other: Self) -> Self::Output {
69 FFTComplex { re: self.re * other.re - self.im * other.im,
70 im: self.im * other.re + self.re * other.im }
74 impl MulAssign for FFTComplex {
75 fn mul_assign(&mut self, other: Self) {
76 let re = self.re * other.re - self.im * other.im;
77 let im = self.im * other.re + self.re * other.im;
83 impl fmt::Display for FFTComplex {
84 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
85 write!(f, "({}, {})", self.re, self.im)
89 pub const FFTC_ZERO: FFTComplex = FFTComplex { re: 0.0, im: 0.0 };
91 #[derive(Debug,Clone,Copy,PartialEq)]
99 table: Vec<FFTComplex>,
107 fn do_fft_inplace_ct(&mut self, data: &mut [FFTComplex], bits: u32, forward: bool) {
108 if bits == 0 { return; }
110 let sum01 = data[0] + data[1];
111 let dif01 = data[0] - data[1];
117 let sum01 = data[0] + data[1];
118 let dif01 = data[0] - data[1];
119 let sum23 = data[2] + data[3];
120 let dif23 = data[2] - data[3];
122 data[0] = sum01 + sum23;
123 data[1] = dif01 - dif23.rotate();
124 data[2] = sum01 - sum23;
125 data[3] = dif01 + dif23.rotate();
127 data[0] = sum01 + sum23;
128 data[1] = dif01 + dif23.rotate();
129 data[2] = sum01 - sum23;
130 data[3] = dif01 - dif23.rotate();
135 let hsize = (1 << (bits - 1)) as usize;
136 self.do_fft_inplace_ct(&mut data[0..hsize], bits - 1, forward);
137 self.do_fft_inplace_ct(&mut data[hsize..], bits - 1, forward);
148 let o = data[k + hsize] * self.table[offs + k];
150 data[k + hsize] = e - o;
155 let o = data[k + hsize] * !self.table[offs + k];
157 data[k + hsize] = e - o;
162 fn do_fft_inplace_splitradix(&mut self, data: &mut [FFTComplex], bits: u32, forward: bool) {
163 if bits == 0 { return; }
165 let sum01 = data[0] + data[1];
166 let dif01 = data[0] - data[1];
172 let sum01 = data[0] + data[2];
173 let dif01 = data[0] - data[2];
174 let sum23 = data[1] + data[3];
175 let dif23 = data[1] - data[3];
177 data[0] = sum01 + sum23;
178 data[1] = dif01 - dif23.rotate();
179 data[2] = sum01 - sum23;
180 data[3] = dif01 + dif23.rotate();
182 data[0] = sum01 + sum23;
183 data[1] = dif01 + dif23.rotate();
184 data[2] = sum01 - sum23;
185 data[3] = dif01 - dif23.rotate();
189 let qsize = (1 << (bits - 2)) as usize;
190 let hsize = (1 << (bits - 1)) as usize;
191 let q3size = qsize + hsize;
193 self.do_fft_inplace_splitradix(&mut data[0 ..hsize], bits - 1, forward);
194 self.do_fft_inplace_splitradix(&mut data[hsize ..q3size], bits - 2, forward);
195 self.do_fft_inplace_splitradix(&mut data[q3size..], bits - 2, forward);
199 let t3 = data[0 + hsize] + data[0 + q3size];
200 let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
202 let e2 = data[0 + qsize];
204 data[0 + qsize] = e2 - t4;
205 data[0 + hsize] = e1 - t3;
206 data[0 + q3size] = e2 + t4;
209 let t1 = self.table[off + k * 2 + 0] * data[k + hsize];
210 let t2 = self.table[off + k * 2 + 1] * data[k + q3size];
212 let t4 = (t1 - t2).rotate();
214 let e2 = data[k + qsize];
216 data[k + qsize] = e2 - t4;
217 data[k + hsize] = e1 - t3;
218 data[k + qsize * 3] = e2 + t4;
222 let t3 = data[0 + hsize] + data[0 + q3size];
223 let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
225 let e2 = data[0 + qsize];
227 data[0 + qsize] = e2 + t4;
228 data[0 + hsize] = e1 - t3;
229 data[0 + q3size] = e2 - t4;
232 let t1 = !self.table[off + k * 2 + 0] * data[k + hsize];
233 let t2 = !self.table[off + k * 2 + 1] * data[k + q3size];
235 let t4 = (t1 - t2).rotate();
237 let e2 = data[k + qsize];
239 data[k + qsize] = e2 + t4;
240 data[k + hsize] = e1 - t3;
241 data[k + qsize * 3] = e2 - t4;
246 pub fn do_fft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex], forward: bool) {
249 let base = if forward { -consts::PI * 2.0 / (src.len() as f32) }
250 else { consts::PI * 2.0 / (src.len() as f32) };
251 for k in 0..src.len() {
252 let mut sum = FFTC_ZERO;
253 for n in 0..src.len() {
254 let w = FFTComplex::exp(base * ((n * k) as f32));
260 FFTMode::CooleyTukey => {
261 let bits = self.bits;
262 for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
263 self.do_fft_inplace_ct(dst, bits, forward);
265 FFTMode::SplitRadix => {
266 let bits = self.bits;
267 for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
268 self.do_fft_inplace_splitradix(dst, bits, forward);
273 pub fn do_fft_inplace(&mut self, data: &mut [FFTComplex], forward: bool) {
274 for idx in 0..self.swaps.len() {
275 let nidx = self.swaps[idx];
278 data[nidx] = data[idx];
284 let size = (1 << self.bits) as usize;
285 let base = if forward { -consts::PI * 2.0 / (size as f32) }
286 else { consts::PI * 2.0 / (size as f32) };
287 let mut res: Vec<FFTComplex> = Vec::with_capacity(size);
289 let mut sum = FFTC_ZERO;
291 let w = FFTComplex::exp(base * ((n * k) as f32));
300 FFTMode::CooleyTukey => {
301 let bits = self.bits;
302 self.do_fft_inplace_ct(data, bits, forward);
304 FFTMode::SplitRadix => {
305 let bits = self.bits;
306 self.do_fft_inplace_splitradix(data, bits, forward);
312 pub struct FFTBuilder {
315 fn reverse_bits(inval: u32) -> u32 {
316 const REV_TAB: [u8; 16] = [
317 0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
318 0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
324 ret = (ret << 4) | (REV_TAB[(val & 0xF) as usize] as u32);
330 fn swp_idx(idx: usize, bits: u32) -> usize {
331 let s = reverse_bits(idx as u32) as usize;
335 fn gen_sr_perms(swaps: &mut [usize], size: usize) {
336 if size <= 4 { return; }
337 let mut evec: Vec<usize> = Vec::with_capacity(size / 2);
338 let mut ovec1: Vec<usize> = Vec::with_capacity(size / 4);
339 let mut ovec2: Vec<usize> = Vec::with_capacity(size / 4);
341 evec.push (swaps[k * 4 + 0]);
342 ovec1.push(swaps[k * 4 + 1]);
343 evec.push (swaps[k * 4 + 2]);
344 ovec2.push(swaps[k * 4 + 3]);
346 for k in 0..size/2 { swaps[k] = evec[k]; }
347 for k in 0..size/4 { swaps[k + size/2] = ovec1[k]; }
348 for k in 0..size/4 { swaps[k + 3*size/4] = ovec2[k]; }
349 gen_sr_perms(&mut swaps[0..size/2], size/2);
350 gen_sr_perms(&mut swaps[size/2..3*size/4], size/4);
351 gen_sr_perms(&mut swaps[3*size/4..], size/4);
354 fn gen_swaps_for_perm(swaps: &mut Vec<usize>, perms: &Vec<usize>) {
355 let mut idx_arr: Vec<usize> = Vec::with_capacity(perms.len());
356 for i in 0..perms.len() { idx_arr.push(i); }
357 let mut run_size = 0;
359 for idx in 0..perms.len() {
360 if perms[idx] == idx_arr[idx] {
361 if run_size == 0 { run_pos = idx; }
364 for i in 0..run_size {
365 swaps.push(run_pos + i);
368 let mut spos = idx + 1;
369 while idx_arr[spos] != perms[idx] { spos += 1; }
370 idx_arr[spos] = idx_arr[idx];
371 idx_arr[idx] = perms[idx];
378 pub fn new_fft(mode: FFTMode, size: usize) -> FFT {
379 let mut swaps: Vec<usize>;
380 let mut perms: Vec<usize>;
381 let mut table: Vec<FFTComplex>;
382 let bits = 31 - (size as u32).leading_zeros();
389 FFTMode::CooleyTukey => {
390 perms = Vec::with_capacity(size);
392 perms.push(swp_idx(i, bits));
394 swaps = Vec::with_capacity(size);
395 table = Vec::with_capacity(size);
396 for _ in 0..4 { table.push(FFTC_ZERO); }
397 for b in 3..(bits+1) {
398 let hsize = (1 << (b - 1)) as usize;
399 let base = -consts::PI / (hsize as f32);
401 table.push(FFTComplex::exp(base * (k as f32)));
405 FFTMode::SplitRadix => {
406 perms = Vec::with_capacity(size);
410 gen_sr_perms(perms.as_mut_slice(), 1 << bits);
411 swaps = Vec::with_capacity(size);
412 table = Vec::with_capacity(size);
413 for _ in 0..4 { table.push(FFTC_ZERO); }
414 for b in 3..(bits+1) {
415 let qsize = (1 << (b - 2)) as usize;
416 let base = -consts::PI / ((qsize * 2) as f32);
418 table.push(FFTComplex::exp(base * ((k * 1) as f32)));
419 table.push(FFTComplex::exp(base * ((k * 3) as f32)));
424 gen_swaps_for_perm(&mut swaps, &perms);
425 FFT { mode: mode, swaps: swaps, perms: perms, bits: bits, table: table }
430 table: Vec<FFTComplex>,
437 fn crossadd(a: &FFTComplex, b: &FFTComplex) -> FFTComplex {
438 FFTComplex { re: a.re + b.re, im: a.im - b.im }
442 pub fn do_rdft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
443 dst.copy_from_slice(src);
444 self.do_rdft_inplace(dst);
446 pub fn do_rdft_inplace(&mut self, buf: &mut [FFTComplex]) {
448 for n in 0..self.size/2 {
449 let in0 = buf[n + 1];
450 let in1 = buf[self.size - n - 1];
452 let t0 = crossadd(&in0, &in1);
453 let t1 = FFTComplex { re: in1.im + in0.im, im: in1.re - in0.re };
454 let tab = self.table[n];
455 let t2 = FFTComplex { re: t1.im * tab.im + t1.re * tab.re, im: t1.im * tab.re - t1.re * tab.im };
457 buf[n + 1] = FFTComplex { re: t0.im - t2.im, im: t0.re - t2.re }; // (t0 - t2).conj().rotate()
458 buf[self.size - n - 1] = (t0 + t2).rotate();
465 self.fft.do_fft_inplace(buf, self.fwd_fft);
467 for n in 0..self.size/2 {
468 let in0 = buf[n + 1];
469 let in1 = buf[self.size - n - 1];
471 let t0 = crossadd(&in0, &in1).scale(0.5);
472 let t1 = FFTComplex { re: in0.im + in1.im, im: in0.re - in1.re };
473 let t2 = t1 * self.table[n];
475 buf[n + 1] = crossadd(&t0, &t2);
476 buf[self.size - n - 1] = FFTComplex { re: t0.re - t2.re, im: -(t0.im + t2.im) };
483 for n in 0..self.size {
484 buf[n] = FFTComplex{ re: buf[n].im, im: buf[n].re };
490 pub struct RDFTBuilder {
494 pub fn new_rdft(mode: FFTMode, size: usize, forward: bool, forward_fft: bool) -> RDFT {
495 let mut table: Vec<FFTComplex> = Vec::with_capacity(size / 4);
496 let (base, scale) = if forward { (consts::PI / (size as f32), 0.5) } else { (-consts::PI / (size as f32), 1.0) };
498 table.push(FFTComplex::exp(base * ((i + 1) as f32)).scale(scale));
500 let fft = FFTBuilder::new_fft(mode, size);
501 RDFT { table, fft, size, fwd: forward, fwd_fft: forward_fft }
512 let mut fin: [FFTComplex; 128] = [FFTC_ZERO; 128];
513 let mut fout1: [FFTComplex; 128] = [FFTC_ZERO; 128];
514 let mut fout2: [FFTComplex; 128] = [FFTC_ZERO; 128];
515 let mut fout3: [FFTComplex; 128] = [FFTC_ZERO; 128];
516 let mut fft1 = FFTBuilder::new_fft(FFTMode::Matrix, fin.len());
517 let mut fft2 = FFTBuilder::new_fft(FFTMode::CooleyTukey, fin.len());
518 let mut fft3 = FFTBuilder::new_fft(FFTMode::SplitRadix, fin.len());
519 let mut seed: u32 = 42;
520 for i in 0..fin.len() {
521 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
522 let val = (seed >> 16) as i16;
523 fin[i].re = (val as f32) / 256.0;
524 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
525 let val = (seed >> 16) as i16;
526 fin[i].im = (val as f32) / 256.0;
528 fft1.do_fft(&fin, &mut fout1, true);
529 fft2.do_fft(&fin, &mut fout2, true);
530 fft3.do_fft(&fin, &mut fout3, true);
532 for i in 0..fin.len() {
533 assert!((fout1[i].re - fout2[i].re).abs() < 1.0);
534 assert!((fout1[i].im - fout2[i].im).abs() < 1.0);
535 assert!((fout1[i].re - fout3[i].re).abs() < 1.0);
536 assert!((fout1[i].im - fout3[i].im).abs() < 1.0);
538 fft1.do_fft_inplace(&mut fout1, false);
539 fft2.do_fft_inplace(&mut fout2, false);
540 fft3.do_fft_inplace(&mut fout3, false);
542 let sc = 1.0 / (fin.len() as f32);
543 for i in 0..fin.len() {
544 assert!((fin[i].re - fout1[i].re * sc).abs() < 1.0);
545 assert!((fin[i].im - fout1[i].im * sc).abs() < 1.0);
546 assert!((fout1[i].re - fout2[i].re).abs() < 1.0);
547 assert!((fout1[i].im - fout2[i].im).abs() < 1.0);
548 assert!((fout1[i].re - fout3[i].re).abs() < 1.0);
549 assert!((fout1[i].im - fout3[i].im).abs() < 1.0);
555 let mut fin: [FFTComplex; 128] = [FFTC_ZERO; 128];
556 let mut fout1: [FFTComplex; 128] = [FFTC_ZERO; 128];
557 let mut rdft = RDFTBuilder::new_rdft(FFTMode::SplitRadix, fin.len(), true, true);
558 let mut seed: u32 = 42;
559 for i in 0..fin.len() {
560 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
561 let val = (seed >> 16) as i16;
562 fin[i].re = (val as f32) / 256.0;
563 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
564 let val = (seed >> 16) as i16;
565 fin[i].im = (val as f32) / 256.0;
567 rdft.do_rdft(&fin, &mut fout1);
568 let mut irdft = RDFTBuilder::new_rdft(FFTMode::SplitRadix, fin.len(), false, true);
569 irdft.do_rdft_inplace(&mut fout1);
571 for i in 0..fin.len() {
572 let tst = fout1[i].scale(0.5/(fout1.len() as f32));
573 assert!((tst.re - fin[i].re).abs() < 1.0);
574 assert!((tst.im - fin[i].im).abs() < 1.0);