codec-support: add division support for FFTComplex
[nihav.git] / nihav-codec-support / src / dsp / fft.rs
CommitLineData
4e034a32 1//! FFT and RDFT implementation.
3f3b5b23 2use std::f32::{self, consts};
03d914bb 3use std::ops::{Not, Neg, Add, AddAssign, Sub, SubAssign, Mul, MulAssign, Div};
3f3b5b23
KS
4use std::fmt;
5
4e034a32 6/// Complex number.
b01971f2 7#[repr(C)]
3f3b5b23
KS
8#[derive(Debug,Clone,Copy,PartialEq)]
9pub struct FFTComplex {
4e034a32 10 /// Real part of the numner.
3f3b5b23 11 pub re: f32,
4e034a32 12 /// Complex part of the number.
3f3b5b23
KS
13 pub im: f32,
14}
15
16impl FFTComplex {
4e034a32 17 /// Calculates `exp(i * val)`.
3f3b5b23
KS
18 pub fn exp(val: f32) -> Self {
19 FFTComplex { re: val.cos(), im: val.sin() }
20 }
4e034a32 21 /// Returns `-Im + i * Re`.
3f3b5b23
KS
22 pub fn rotate(self) -> Self {
23 FFTComplex { re: -self.im, im: self.re }
24 }
4e034a32 25 /// Multiplies complex number by scalar.
3f3b5b23
KS
26 pub fn scale(self, scale: f32) -> Self {
27 FFTComplex { re: self.re * scale, im: self.im * scale }
28 }
03d914bb
KS
29 /// Returns squared modulus value of the complex number.
30 pub fn sq_modulus(self) -> f32 {
31 self.re * self.re + self.im * self.im
32 }
33 /// Returns reciprocal of the complex number.
34 pub fn reciprocal(self) -> Self {
35 !self.scale(self.sq_modulus())
36 }
3f3b5b23
KS
37}
38
39impl Neg for FFTComplex {
40 type Output = FFTComplex;
41 fn neg(self) -> Self::Output {
42 FFTComplex { re: -self.re, im: -self.im }
43 }
44}
45
46impl Not for FFTComplex {
47 type Output = FFTComplex;
48 fn not(self) -> Self::Output {
49 FFTComplex { re: self.re, im: -self.im }
50 }
51}
52
53impl Add for FFTComplex {
54 type Output = FFTComplex;
55 fn add(self, other: Self) -> Self::Output {
56 FFTComplex { re: self.re + other.re, im: self.im + other.im }
57 }
58}
59
60impl AddAssign for FFTComplex {
61 fn add_assign(&mut self, other: Self) {
62 self.re += other.re;
63 self.im += other.im;
64 }
65}
66
67impl Sub for FFTComplex {
68 type Output = FFTComplex;
69 fn sub(self, other: Self) -> Self::Output {
70 FFTComplex { re: self.re - other.re, im: self.im - other.im }
71 }
72}
73
74impl SubAssign for FFTComplex {
75 fn sub_assign(&mut self, other: Self) {
76 self.re -= other.re;
77 self.im -= other.im;
78 }
79}
80
81impl Mul for FFTComplex {
82 type Output = FFTComplex;
83 fn mul(self, other: Self) -> Self::Output {
84 FFTComplex { re: self.re * other.re - self.im * other.im,
85 im: self.im * other.re + self.re * other.im }
86 }
87}
88
89impl MulAssign for FFTComplex {
90 fn mul_assign(&mut self, other: Self) {
91 let re = self.re * other.re - self.im * other.im;
92 let im = self.im * other.re + self.re * other.im;
93 self.re = re;
94 self.im = im;
95 }
96}
97
03d914bb
KS
98impl Div for FFTComplex {
99 type Output = FFTComplex;
100 fn div(self, other: Self) -> Self::Output {
101 self * other.reciprocal()
102 }
103}
104
3f3b5b23
KS
105impl fmt::Display for FFTComplex {
106 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
107 write!(f, "({}, {})", self.re, self.im)
108 }
109}
110
4e034a32 111/// Complex number with zero value.
3f3b5b23
KS
112pub const FFTC_ZERO: FFTComplex = FFTComplex { re: 0.0, im: 0.0 };
113
4e034a32 114/// Calculates forward or inverse FFT in the straightforward way.
9e78289c
KS
115pub fn generic_fft(data: &mut [FFTComplex], forward: bool) {
116 let mut tmp = Vec::with_capacity(data.len());
117 tmp.resize(data.len(), FFTC_ZERO);
118 let base = if forward { -consts::PI * 2.0 / (data.len() as f32) }
119 else { consts::PI * 2.0 / (data.len() as f32) };
120 for k in 0..data.len() {
121 let mut sum = FFTC_ZERO;
122 for n in 0..data.len() {
123 let w = FFTComplex::exp(base * ((n * k) as f32));
124 sum += data[n] * w;
125 }
126 tmp[k] = sum;
127 }
128 for k in 0..data.len() {
129 data[k] = tmp[k];
130 }
3f3b5b23
KS
131}
132
9e78289c
KS
133struct FFTData {
134 table: Vec<FFTComplex>,
135 tmp: Vec<FFTComplex>,
136 twiddle: Vec<FFTComplex>,
137 size: usize,
138 step: usize,
139 div: usize,
3f3b5b23
KS
140}
141
9e78289c
KS
142struct FFTGeneric {}
143
144const FFT3_CONST: f32 = 0.86602540378443864677;
145const FFT5_CONST1: FFTComplex = FFTComplex { re: 0.80901699437494742410, im: 0.58778525229247312915 };
146const FFT5_CONST2: FFTComplex = FFTComplex { re: 0.30901699437494742411, im: 0.95105651629515357211 };
147
148fn twiddle5(a: FFTComplex, b: FFTComplex, c: FFTComplex) -> (FFTComplex, FFTComplex) {
149 let re = a.re * c.re;
150 let im = a.im * c.re;
151 let diffre = b.im * c.im;
152 let diffim = b.re * c.im;
153
154 (FFTComplex { re: re - diffre, im: im + diffim }, FFTComplex { re: re + diffre, im: im - diffim })
155}
156
157impl FFTGeneric {
158 fn new_data(size: usize, forward: bool) -> FFTData {
159 let mut table: Vec<FFTComplex> = Vec::with_capacity(size * size);
160 table.resize(size * size, FFTC_ZERO);
161 let base = consts::PI * 2.0 / (size as f32);
162 if forward {
163 for n in 0..size {
164 for k in 0..size {
165 table[n * size + k] = FFTComplex::exp(-base * ((n * k) as f32));
166 }
167 }
168 } else {
169 for n in 0..size {
170 for k in 0..size {
171 table[n * size + k] = FFTComplex::exp( base * ((n * k) as f32));
172 }
173 }
174 }
175 let mut tmp = Vec::with_capacity(size);
176 tmp.resize(size, FFTC_ZERO);
177 FFTData { table, tmp, twiddle: Vec::new(), size, step: 0, div: 0 }
178 }
179 fn fft(tbl: &mut FFTData, size: usize, data: &mut [FFTComplex], step: usize) {
180 if size == 3 {
181 let s0 = data[step * 0];
182 let s1 = data[step * 1];
183 let s2 = data[step * 2];
184 let t0 = s1 + s2;
185 data[step * 0] += t0;
186 let t1 = s0 - t0.scale(0.5);
187 let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
188 data[step * 1] = t1 + t2;
189 data[step * 2] = t1 - t2;
190 return;
191 }
192 if size == 5 {
193 let s0 = data[step * 0];
194 let s1 = data[step * 1];
195 let s2 = data[step * 2];
196 let s3 = data[step * 3];
197 let s4 = data[step * 4];
198
199 let t0 = s1 + s4;
200 let t1 = s1 - s4;
201 let t2 = s2 + s3;
202 let t3 = s2 - s3;
203 let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
204 let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
205 let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
206 let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
207
208 data[step * 0] = s0 + t0 + t2;
209 data[step * 1] = s0 + t5 - t6;
210 data[step * 2] = s0 - t8 + ta;
211 data[step * 3] = s0 - t9 + tb;
212 data[step * 4] = s0 + t4 - t7;
213 return;
214 }
215 for k in 0..tbl.size {
216 tbl.tmp[k] = FFTC_ZERO;
217 for n in 0..tbl.size {
218 tbl.tmp[k] += data[n * step] * tbl.table[k * tbl.size + n];
219 }
220 }
221 for n in 0..tbl.size {
222 data[n * step] = tbl.tmp[n];
223 }
224 }
225 fn ifft(tbl: &mut FFTData, size: usize, data: &mut [FFTComplex], step: usize) {
226 if size == 3 {
227 let s0 = data[step * 0];
228 let s1 = data[step * 1];
229 let s2 = data[step * 2];
230 let t0 = s1 + s2;
231 data[step * 0] += t0;
232 let t1 = s0 - t0.scale(0.5);
233 let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
234 data[step * 1] = t1 - t2;
235 data[step * 2] = t1 + t2;
236 return;
237 }
238 if size == 5 {
239 let s0 = data[step * 0];
240 let s1 = data[step * 1];
241 let s2 = data[step * 2];
242 let s3 = data[step * 3];
243 let s4 = data[step * 4];
244
245 let t0 = s1 + s4;
246 let t1 = s1 - s4;
247 let t2 = s2 + s3;
248 let t3 = s2 - s3;
249 let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
250 let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
251 let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
252 let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
253
254 data[step * 0] = s0 + t0 + t2;
255 data[step * 1] = s0 + t4 - t7;
256 data[step * 2] = s0 - t9 + tb;
257 data[step * 3] = s0 - t8 + ta;
258 data[step * 4] = s0 + t5 - t6;
259 return;
260 }
261 Self::fft(tbl, size, data, step);
262 }
263}
264
265struct FFTSplitRadix {}
266
267impl FFTSplitRadix {
268 fn new_data(bits: u8, _forward: bool) -> FFTData {
269 let size = 1 << bits;
270 let mut table = Vec::with_capacity(size);
271 for _ in 0..4 { table.push(FFTC_ZERO); }
fdb4b2fb 272 for b in 3..=bits {
9e78289c
KS
273 let qsize = (1 << (b - 2)) as usize;
274 let base = -consts::PI / ((qsize * 2) as f32);
275 for k in 0..qsize {
276 table.push(FFTComplex::exp(base * ((k * 1) as f32)));
277 table.push(FFTComplex::exp(base * ((k * 3) as f32)));
278 }
279 }
280 FFTData { table, tmp: Vec::new(), twiddle: Vec::new(), size, step: 0, div: 0 }
281 }
282 fn fft(fftdata: &mut FFTData, bits: u8, data: &mut [FFTComplex]) {
3f3b5b23
KS
283 if bits == 0 { return; }
284 if bits == 1 {
285 let sum01 = data[0] + data[1];
286 let dif01 = data[0] - data[1];
287 data[0] = sum01;
288 data[1] = dif01;
289 return;
290 }
291 if bits == 2 {
9e78289c
KS
292 let sum01 = data[0] + data[2];
293 let dif01 = data[0] - data[2];
294 let sum23 = data[1] + data[3];
295 let dif23 = data[1] - data[3];
296 data[0] = sum01 + sum23;
297 data[1] = dif01 - dif23.rotate();
298 data[2] = sum01 - sum23;
299 data[3] = dif01 + dif23.rotate();
3f3b5b23
KS
300 return;
301 }
9e78289c 302 let qsize = (1 << (bits - 2)) as usize;
3f3b5b23 303 let hsize = (1 << (bits - 1)) as usize;
9e78289c
KS
304 let q3size = qsize + hsize;
305
306 Self::fft(fftdata, bits - 1, &mut data[0 ..hsize]);
307 Self::fft(fftdata, bits - 2, &mut data[hsize ..q3size]);
308 Self::fft(fftdata, bits - 2, &mut data[q3size..]);
309 let off = hsize;
3f3b5b23 310 {
9e78289c
KS
311 let t3 = data[0 + hsize] + data[0 + q3size];
312 let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
313 let e1 = data[0];
314 let e2 = data[0 + qsize];
315 data[0] = e1 + t3;
316 data[0 + qsize] = e2 - t4;
317 data[0 + hsize] = e1 - t3;
318 data[0 + q3size] = e2 + t4;
3f3b5b23 319 }
9e78289c
KS
320 for k in 1..qsize {
321 let t1 = fftdata.table[off + k * 2 + 0] * data[k + hsize];
322 let t2 = fftdata.table[off + k * 2 + 1] * data[k + q3size];
323 let t3 = t1 + t2;
324 let t4 = (t1 - t2).rotate();
325 let e1 = data[k];
326 let e2 = data[k + qsize];
327 data[k] = e1 + t3;
328 data[k + qsize] = e2 - t4;
329 data[k + hsize] = e1 - t3;
330 data[k + qsize * 3] = e2 + t4;
3f3b5b23
KS
331 }
332 }
9e78289c 333 fn ifft(fftdata: &mut FFTData, bits: u8, data: &mut [FFTComplex]) {
3f3b5b23
KS
334 if bits == 0 { return; }
335 if bits == 1 {
336 let sum01 = data[0] + data[1];
337 let dif01 = data[0] - data[1];
338 data[0] = sum01;
339 data[1] = dif01;
340 return;
341 }
342 if bits == 2 {
343 let sum01 = data[0] + data[2];
344 let dif01 = data[0] - data[2];
345 let sum23 = data[1] + data[3];
346 let dif23 = data[1] - data[3];
9e78289c
KS
347 data[0] = sum01 + sum23;
348 data[1] = dif01 + dif23.rotate();
349 data[2] = sum01 - sum23;
350 data[3] = dif01 - dif23.rotate();
3f3b5b23
KS
351 return;
352 }
353 let qsize = (1 << (bits - 2)) as usize;
354 let hsize = (1 << (bits - 1)) as usize;
355 let q3size = qsize + hsize;
356
9e78289c
KS
357 Self::ifft(fftdata, bits - 1, &mut data[0 ..hsize]);
358 Self::ifft(fftdata, bits - 2, &mut data[hsize ..q3size]);
359 Self::ifft(fftdata, bits - 2, &mut data[q3size..]);
3f3b5b23 360 let off = hsize;
9e78289c
KS
361 {
362 let t3 = data[0 + hsize] + data[0 + q3size];
363 let t4 = (data[0 + hsize] - data[0 + q3size]).rotate();
364 let e1 = data[0];
365 let e2 = data[0 + qsize];
366 data[0] = e1 + t3;
367 data[0 + qsize] = e2 + t4;
368 data[0 + hsize] = e1 - t3;
369 data[0 + q3size] = e2 - t4;
370 }
371 for k in 1..qsize {
372 let t1 = !fftdata.table[off + k * 2 + 0] * data[k + hsize];
373 let t2 = !fftdata.table[off + k * 2 + 1] * data[k + q3size];
374 let t3 = t1 + t2;
375 let t4 = (t1 - t2).rotate();
376 let e1 = data[k];
377 let e2 = data[k + qsize];
378 data[k] = e1 + t3;
379 data[k + qsize] = e2 + t4;
380 data[k + hsize] = e1 - t3;
381 data[k + qsize * 3] = e2 - t4;
3f3b5b23
KS
382 }
383 }
9e78289c
KS
384}
385
386struct FFT15 {}
3f3b5b23 387
9e78289c
KS
388const FFT15_INSWAP: [usize; 20] = [ 0, 5, 10, 42, 3, 8, 13, 42, 6, 11, 1, 42, 9, 14, 4, 42, 12, 2, 7, 42 ];
389const FFT15_OUTSWAP: [usize; 20] = [ 0, 10, 5, 42, 6, 1, 11, 42, 12, 7, 2, 42, 3, 13, 8, 42, 9, 4, 14, 42 ];
390
391impl FFT15 {
392 fn new_data(size: usize, _forward: bool) -> FFTData {
393 FFTData { table: Vec::new(), tmp: Vec::new(), twiddle: Vec::new(), size, step: 0, div: 0 }
394 }
395 fn fft3(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
396 let s0 = src[0];
397 let s1 = src[1];
398 let s2 = src[2];
399
400 let t0 = s1 + s2;
401 let t1 = s0 - t0.scale(0.5);
402 let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
403
404 dst[FFT15_OUTSWAP[n * 4 + 0] * step] = s0 + t0;
405 dst[FFT15_OUTSWAP[n * 4 + 1] * step] = t1 + t2;
406 dst[FFT15_OUTSWAP[n * 4 + 2] * step] = t1 - t2;
407 }
408 fn ifft3(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
409 let s0 = src[0];
410 let s1 = src[1];
411 let s2 = src[2];
412
413 let t0 = s1 + s2;
414 let t1 = s0 - t0.scale(0.5);
415 let t2 = (s2 - s1).rotate().scale(FFT3_CONST);
416
417 dst[FFT15_OUTSWAP[n * 4 + 0] * step] = s0 + t0;
418 dst[FFT15_OUTSWAP[n * 4 + 1] * step] = t1 - t2;
419 dst[FFT15_OUTSWAP[n * 4 + 2] * step] = t1 + t2;
420 }
421 fn fft5(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
422 let s0 = src[FFT15_INSWAP[n + 0 * 4] * step];
423 let s1 = src[FFT15_INSWAP[n + 1 * 4] * step];
424 let s2 = src[FFT15_INSWAP[n + 2 * 4] * step];
425 let s3 = src[FFT15_INSWAP[n + 3 * 4] * step];
426 let s4 = src[FFT15_INSWAP[n + 4 * 4] * step];
427
428 let t0 = s1 + s4;
429 let t1 = s1 - s4;
430 let t2 = s2 + s3;
431 let t3 = s2 - s3;
432 let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
433 let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
434 let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
435 let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
436
437 dst[0 * 3] = s0 + t0 + t2;
438 dst[1 * 3] = s0 + t5 - t6;
439 dst[2 * 3] = s0 - t8 + ta;
440 dst[3 * 3] = s0 - t9 + tb;
441 dst[4 * 3] = s0 + t4 - t7;
442 }
443 fn ifft5(dst: &mut [FFTComplex], src: &[FFTComplex], step: usize, n: usize) {
444 let s0 = src[FFT15_INSWAP[n + 0 * 4] * step];
445 let s1 = src[FFT15_INSWAP[n + 1 * 4] * step];
446 let s2 = src[FFT15_INSWAP[n + 2 * 4] * step];
447 let s3 = src[FFT15_INSWAP[n + 3 * 4] * step];
448 let s4 = src[FFT15_INSWAP[n + 4 * 4] * step];
449
450 let t0 = s1 + s4;
451 let t1 = s1 - s4;
452 let t2 = s2 + s3;
453 let t3 = s2 - s3;
454 let (t4, t5) = twiddle5(t0, t1, FFT5_CONST2);
455 let (t6, t7) = twiddle5(t2, t3, FFT5_CONST1);
456 let (t8, t9) = twiddle5(t0, t1, FFT5_CONST1);
457 let (ta, tb) = twiddle5(t2, t3, FFT5_CONST2);
458
459 dst[0 * 3] = s0 + t0 + t2;
460 dst[1 * 3] = s0 + t4 - t7;
461 dst[2 * 3] = s0 - t9 + tb;
462 dst[3 * 3] = s0 - t8 + ta;
463 dst[4 * 3] = s0 + t5 - t6;
464 }
465 fn fft(_fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
466 let mut tmp = [FFTC_ZERO; 15];
467 for n in 0..3 {
468 Self::fft5(&mut tmp[n..], data, step, n);
469 }
470 for n in 0..5 {
471 Self::fft3(data, &tmp[n * 3..][..3], step, n);
472 }
473 }
474 fn ifft(_fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
475 let mut tmp = [FFTC_ZERO; 15];
476 for n in 0..3 {
477 Self::ifft5(&mut tmp[n..], data, step, n);
478 }
479 for n in 0..5 {
480 Self::ifft3(data, &tmp[n * 3..][..3], step, n);
481 }
482 }
483}
484
485
486enum FFTMode {
487 Generic(usize),
488 SplitRadix(u8),
489 Prime15,
490}
491
492impl FFTMode {
493 fn permute(&self, perms: &mut [usize]) {
494 match *self {
495 FFTMode::Generic(_) => {},
496 FFTMode::SplitRadix(bits) => {
497 let div = perms.len() >> bits;
498 gen_sr_perms(perms, 1 << bits);
499 if div > 1 {
500 for i in 0..(1 << bits) {
501 perms[i] *= div;
502 }
503 for i in 1..div {
504 for j in 0..(1 << bits) {
505 perms[(i << bits) + j] = perms[j] + i;
3f3b5b23 506 }
3f3b5b23 507 }
9e78289c
KS
508 }
509 },
510 FFTMode::Prime15 => {},
3f3b5b23
KS
511 };
512 }
9e78289c
KS
513 fn do_fft(&self, fftdata: &mut FFTData, data: &mut [FFTComplex]) {
514 match *self {
515 FFTMode::Generic(size) => FFTGeneric::fft(fftdata, size, data, 1),
516 FFTMode::SplitRadix(bits) => FFTSplitRadix::fft(fftdata, bits, data),
517 FFTMode::Prime15 => FFT15::fft(fftdata, data, 1),
518 };
519 }
520 fn do_fft2(&self, fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
521 match *self {
522 FFTMode::Generic(size) => FFTGeneric::fft(fftdata, size, data, step),
523 FFTMode::SplitRadix(_) => unreachable!(),
524 FFTMode::Prime15 => FFT15::fft(fftdata, data, step),
525 };
526 }
527 fn do_ifft(&self, fftdata: &mut FFTData, data: &mut [FFTComplex]) {
528 match *self {
529 FFTMode::Generic(size) => FFTGeneric::ifft(fftdata, size, data, 1),
530 FFTMode::SplitRadix(bits) => FFTSplitRadix::ifft(fftdata, bits, data),
531 FFTMode::Prime15 => FFT15::ifft(fftdata, data, 1),
532 };
533 }
534 fn do_ifft2(&self, fftdata: &mut FFTData, data: &mut [FFTComplex], step: usize) {
535 match *self {
536 FFTMode::Generic(size) => FFTGeneric::ifft(fftdata, size, data, step),
537 FFTMode::SplitRadix(_) => unreachable!(),
538 FFTMode::Prime15 => FFT15::ifft(fftdata, data, step),
539 };
540 }
541 fn get_size(&self) -> usize {
542 match *self {
543 FFTMode::Generic(size) => size,
544 FFTMode::SplitRadix(bits) => 1 << bits,
545 FFTMode::Prime15 => 15,
546 }
547 }
548}
3f3b5b23 549
4e034a32 550/// FFT working context.
9e78289c
KS
551pub struct FFT {
552 perms: Vec<usize>,
553 swaps: Vec<usize>,
554 ffts: Vec<(FFTMode, FFTData)>,
555}
556
557impl FFT {
4e034a32 558 /// Calculates Fourier transform.
9e78289c
KS
559 pub fn do_fft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
560 for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
561 self.do_fft_core(dst);
562 }
4e034a32 563 /// Calculates inverse Fourier transform.
9e78289c
KS
564 pub fn do_ifft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
565 for k in 0..src.len() { dst[k] = src[self.perms[k]]; }
566 self.do_ifft_core(dst);
567 }
4e034a32 568 /// Performs inplace FFT.
9e78289c 569 pub fn do_fft_inplace(&mut self, data: &mut [FFTComplex]) {
3f3b5b23
KS
570 for idx in 0..self.swaps.len() {
571 let nidx = self.swaps[idx];
572 if idx != nidx {
e243ceb4 573 data.swap(nidx, idx);
3f3b5b23
KS
574 }
575 }
9e78289c
KS
576 self.do_fft_core(data);
577 }
4e034a32 578 /// Performs inplace inverse FFT.
9e78289c
KS
579 pub fn do_ifft_inplace(&mut self, data: &mut [FFTComplex]) {
580 for idx in 0..self.swaps.len() {
581 let nidx = self.swaps[idx];
582 if idx != nidx {
e243ceb4 583 data.swap(nidx, idx);
9e78289c
KS
584 }
585 }
586 self.do_ifft_core(data);
587 }
588 fn do_fft_core(&mut self, data: &mut [FFTComplex]) {
589 for el in self.ffts.iter_mut() {
590 let (mode, ref mut fftdata) = el;
591 let bsize = mode.get_size();
592 let div = fftdata.div;
593 let step = fftdata.step;
594 if step == 1 {
595 mode.do_fft(fftdata, data);
596 for i in 1..div {
597 mode.do_fft(fftdata, &mut data[i * bsize..]);
598 }
599 } else {
600 mode.do_fft2(fftdata, data, div);
601 let mut toff = bsize;
602 for i in 1..div {
603 for j in 1..bsize {
604 data[i + j * div] *= fftdata.twiddle[toff + j];
3f3b5b23 605 }
9e78289c
KS
606 mode.do_fft2(fftdata, &mut data[i..], div);
607 toff += bsize;
608 }
609 }
610 }
611 }
612 fn do_ifft_core(&mut self, data: &mut [FFTComplex]) {
613 for el in self.ffts.iter_mut() {
614 let (mode, ref mut fftdata) = el;
615 let bsize = mode.get_size();
616 let div = fftdata.div;
617 let step = fftdata.step;
618 if step == 1 {
619 mode.do_ifft(fftdata, data);
620 for i in 1..div {
621 mode.do_ifft(fftdata, &mut data[i * bsize..]);
622 }
623 } else {
624 mode.do_ifft2(fftdata, data, div);
625 let mut toff = bsize;
626 for i in 1..div {
627 for j in 1..bsize {
628 data[i + j * div] *= fftdata.twiddle[toff + j];
3f3b5b23 629 }
9e78289c
KS
630 mode.do_ifft2(fftdata, &mut data[i..], div);
631 toff += bsize;
632 }
633 }
634 }
3f3b5b23
KS
635 }
636}
637
4e034a32
KS
638/// [`FFT`] context creator.
639///
640/// [`FFT`]: ./struct.FFT.html
3f3b5b23
KS
641pub struct FFTBuilder {
642}
643
9e78289c 644/*fn reverse_bits(inval: u32) -> u32 {
3f3b5b23
KS
645 const REV_TAB: [u8; 16] = [
646 0b0000, 0b1000, 0b0100, 0b1100, 0b0010, 0b1010, 0b0110, 0b1110,
647 0b0001, 0b1001, 0b0101, 0b1101, 0b0011, 0b1011, 0b0111, 0b1111,
648 ];
649
650 let mut ret = 0;
651 let mut val = inval;
652 for _ in 0..8 {
653 ret = (ret << 4) | (REV_TAB[(val & 0xF) as usize] as u32);
654 val = val >> 4;
655 }
656 ret
657}
658
659fn swp_idx(idx: usize, bits: u32) -> usize {
660 let s = reverse_bits(idx as u32) as usize;
661 s >> (32 - bits)
9e78289c 662}*/
3f3b5b23
KS
663
664fn gen_sr_perms(swaps: &mut [usize], size: usize) {
665 if size <= 4 { return; }
666 let mut evec: Vec<usize> = Vec::with_capacity(size / 2);
667 let mut ovec1: Vec<usize> = Vec::with_capacity(size / 4);
668 let mut ovec2: Vec<usize> = Vec::with_capacity(size / 4);
669 for k in 0..size/4 {
670 evec.push (swaps[k * 4 + 0]);
671 ovec1.push(swaps[k * 4 + 1]);
672 evec.push (swaps[k * 4 + 2]);
673 ovec2.push(swaps[k * 4 + 3]);
674 }
675 for k in 0..size/2 { swaps[k] = evec[k]; }
676 for k in 0..size/4 { swaps[k + size/2] = ovec1[k]; }
677 for k in 0..size/4 { swaps[k + 3*size/4] = ovec2[k]; }
678 gen_sr_perms(&mut swaps[0..size/2], size/2);
679 gen_sr_perms(&mut swaps[size/2..3*size/4], size/4);
680 gen_sr_perms(&mut swaps[3*size/4..], size/4);
681}
682
fdb4b2fb 683fn gen_swaps_for_perm(swaps: &mut Vec<usize>, perms: &[usize]) {
3f3b5b23
KS
684 let mut idx_arr: Vec<usize> = Vec::with_capacity(perms.len());
685 for i in 0..perms.len() { idx_arr.push(i); }
686 let mut run_size = 0;
687 let mut run_pos = 0;
688 for idx in 0..perms.len() {
689 if perms[idx] == idx_arr[idx] {
690 if run_size == 0 { run_pos = idx; }
691 run_size += 1;
692 } else {
693 for i in 0..run_size {
694 swaps.push(run_pos + i);
695 }
696 run_size = 0;
697 let mut spos = idx + 1;
698 while idx_arr[spos] != perms[idx] { spos += 1; }
699 idx_arr[spos] = idx_arr[idx];
700 idx_arr[idx] = perms[idx];
701 swaps.push(spos);
702 }
703 }
704}
705
706impl FFTBuilder {
9e78289c
KS
707 fn generate_twiddle(data: &mut FFTData, size: usize, cur_size: usize, forward: bool) {
708 if size == cur_size { return; }
709 data.twiddle = Vec::with_capacity(size);
710 let div = size / cur_size;
711 let base = if forward { -2.0 * consts::PI / (size as f32) } else { 2.0 * consts::PI / (size as f32) };
712 for n in 0..div {
713 for k in 0..cur_size {
714 data.twiddle.push(FFTComplex::exp(base * ((k * n) as f32)));
715 }
716 }
717 }
4e034a32 718 /// Constructs a new `FFT` context.
9e78289c
KS
719 pub fn new_fft(size: usize, forward: bool) -> FFT {
720 let mut ffts: Vec<(FFTMode, FFTData)> = Vec::with_capacity(1);
721 let mut perms: Vec<usize> = Vec::with_capacity(size);
722 let mut swaps: Vec<usize> = Vec::with_capacity(size);
723 let mut rem_size = size;
724 if rem_size.trailing_zeros() > 0 {
725 let bits = rem_size.trailing_zeros() as u8;
726 let mut data = FFTSplitRadix::new_data(bits, forward);
727 Self::generate_twiddle(&mut data, size, 1 << bits, forward);
728 data.step = 1;
729 data.div = rem_size >> bits;
730 ffts.push((FFTMode::SplitRadix(bits), data));
731 rem_size >>= bits;
732 }
733 if (rem_size % 15) == 0 {
734 let mut data = FFT15::new_data(size, forward);
735 Self::generate_twiddle(&mut data, size, 15, forward);
736 data.step = size / rem_size;
737 data.div = size / rem_size;
738 ffts.push((FFTMode::Prime15, data));
739 rem_size /= 15;
740 }
741 if rem_size > 1 {
742 let mut data = FFTGeneric::new_data(rem_size, forward);
743 Self::generate_twiddle(&mut data, size, rem_size, forward);
744 data.step = size / rem_size;
745 data.div = size / rem_size;
746 ffts.push((FFTMode::Generic(rem_size), data));
747 }
748
749 for i in 0..size {
750 perms.push(i);
751 }
752 for (mode, _) in ffts.iter().rev() {
753 mode.permute(&mut perms);
754 }
fdb4b2fb 755 gen_swaps_for_perm(&mut swaps, perms.as_slice());
9e78289c
KS
756
757 FFT { perms, swaps, ffts }
3f3b5b23
KS
758 }
759}
760
4e034a32 761/// RDFT working context.
c3528afd
KS
762pub struct RDFT {
763 table: Vec<FFTComplex>,
764 fft: FFT,
765 fwd: bool,
766 size: usize,
b3799375 767 fwd_fft: bool,
c3528afd
KS
768}
769
fdb4b2fb 770fn crossadd(a: FFTComplex, b: FFTComplex) -> FFTComplex {
c3528afd
KS
771 FFTComplex { re: a.re + b.re, im: a.im - b.im }
772}
773
774impl RDFT {
4e034a32 775 /// Calculates RDFT.
c3528afd
KS
776 pub fn do_rdft(&mut self, src: &[FFTComplex], dst: &mut [FFTComplex]) {
777 dst.copy_from_slice(src);
778 self.do_rdft_inplace(dst);
779 }
4e034a32 780 /// Calculates inplace RDFT.
c3528afd
KS
781 pub fn do_rdft_inplace(&mut self, buf: &mut [FFTComplex]) {
782 if !self.fwd {
783 for n in 0..self.size/2 {
784 let in0 = buf[n + 1];
785 let in1 = buf[self.size - n - 1];
786
fdb4b2fb 787 let t0 = crossadd(in0, in1);
c3528afd
KS
788 let t1 = FFTComplex { re: in1.im + in0.im, im: in1.re - in0.re };
789 let tab = self.table[n];
790 let t2 = FFTComplex { re: t1.im * tab.im + t1.re * tab.re, im: t1.im * tab.re - t1.re * tab.im };
791
792 buf[n + 1] = FFTComplex { re: t0.im - t2.im, im: t0.re - t2.re }; // (t0 - t2).conj().rotate()
793 buf[self.size - n - 1] = (t0 + t2).rotate();
794 }
795 let a = buf[0].re;
796 let b = buf[0].im;
797 buf[0].re = a - b;
798 buf[0].im = a + b;
799 }
9e78289c
KS
800 if self.fwd_fft {
801 self.fft.do_fft_inplace(buf);
802 } else {
803 self.fft.do_ifft_inplace(buf);
804 }
c3528afd
KS
805 if self.fwd {
806 for n in 0..self.size/2 {
807 let in0 = buf[n + 1];
808 let in1 = buf[self.size - n - 1];
809
fdb4b2fb 810 let t0 = crossadd(in0, in1).scale(0.5);
c3528afd
KS
811 let t1 = FFTComplex { re: in0.im + in1.im, im: in0.re - in1.re };
812 let t2 = t1 * self.table[n];
813
fdb4b2fb 814 buf[n + 1] = crossadd(t0, t2);
d24468d9 815 buf[self.size - n - 1] = FFTComplex { re: t0.re - t2.re, im: -(t0.im + t2.im) };
c3528afd
KS
816 }
817 let a = buf[0].re;
818 let b = buf[0].im;
819 buf[0].re = a + b;
820 buf[0].im = a - b;
821 } else {
822 for n in 0..self.size {
823 buf[n] = FFTComplex{ re: buf[n].im, im: buf[n].re };
824 }
825 }
826 }
827}
828
4e034a32
KS
829/// [`RDFT`] context creator.
830///
831/// [`RDFT`]: ./struct.FFT.html
c3528afd
KS
832pub struct RDFTBuilder {
833}
834
835impl RDFTBuilder {
4e034a32 836 /// Constructs a new `RDFT` context.
9e78289c 837 pub fn new_rdft(size: usize, forward: bool, forward_fft: bool) -> RDFT {
c3528afd
KS
838 let mut table: Vec<FFTComplex> = Vec::with_capacity(size / 4);
839 let (base, scale) = if forward { (consts::PI / (size as f32), 0.5) } else { (-consts::PI / (size as f32), 1.0) };
840 for i in 0..size/2 {
841 table.push(FFTComplex::exp(base * ((i + 1) as f32)).scale(scale));
842 }
9e78289c 843 let fft = FFTBuilder::new_fft(size, forward_fft);
b3799375 844 RDFT { table, fft, size, fwd: forward, fwd_fft: forward_fft }
c3528afd
KS
845 }
846}
847
3f3b5b23
KS
848
849#[cfg(test)]
850mod test {
851 use super::*;
852
9e78289c
KS
853 fn test_fft(size: usize) {
854 println!("testing FFT {}", size);
855 let mut fin: Vec<FFTComplex> = Vec::with_capacity(size);
856 let mut fout1: Vec<FFTComplex> = Vec::with_capacity(size);
857 let mut fout2: Vec<FFTComplex> = Vec::with_capacity(size);
858 fin.resize(size, FFTC_ZERO);
859 fout1.resize(size, FFTC_ZERO);
860 fout2.resize(size, FFTC_ZERO);
861 let mut fft = FFTBuilder::new_fft(size, true);
3f3b5b23
KS
862 let mut seed: u32 = 42;
863 for i in 0..fin.len() {
864 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
865 let val = (seed >> 16) as i16;
866 fin[i].re = (val as f32) / 256.0;
867 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
868 let val = (seed >> 16) as i16;
869 fin[i].im = (val as f32) / 256.0;
870 }
9e78289c
KS
871 fft.do_fft(&fin, &mut fout1);
872 fout2.copy_from_slice(&fin);
873 generic_fft(&mut fout2, true);
3f3b5b23
KS
874
875 for i in 0..fin.len() {
876 assert!((fout1[i].re - fout2[i].re).abs() < 1.0);
877 assert!((fout1[i].im - fout2[i].im).abs() < 1.0);
3f3b5b23 878 }
9e78289c
KS
879 let mut ifft = FFTBuilder::new_fft(size, false);
880 ifft.do_ifft_inplace(&mut fout1);
881 generic_fft(&mut fout2, false);
3f3b5b23 882
9e78289c 883 let sc = 1.0 / (size as f32);
3f3b5b23
KS
884 for i in 0..fin.len() {
885 assert!((fin[i].re - fout1[i].re * sc).abs() < 1.0);
886 assert!((fin[i].im - fout1[i].im * sc).abs() < 1.0);
9e78289c
KS
887 assert!((fout1[i].re - fout2[i].re).abs() * sc < 1.0);
888 assert!((fout1[i].im - fout2[i].im).abs() * sc < 1.0);
3f3b5b23
KS
889 }
890 }
c3528afd 891
9e78289c
KS
892 #[test]
893 fn test_ffts() {
894 test_fft(3);
895 test_fft(5);
896 test_fft(16);
897 test_fft(15);
898 test_fft(60);
899 test_fft(256);
900 test_fft(240);
901 }
902
c3528afd
KS
903 #[test]
904 fn test_rdft() {
905 let mut fin: [FFTComplex; 128] = [FFTC_ZERO; 128];
906 let mut fout1: [FFTComplex; 128] = [FFTC_ZERO; 128];
9e78289c 907 let mut rdft = RDFTBuilder::new_rdft(fin.len(), true, true);
c3528afd
KS
908 let mut seed: u32 = 42;
909 for i in 0..fin.len() {
910 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
911 let val = (seed >> 16) as i16;
912 fin[i].re = (val as f32) / 256.0;
913 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
914 let val = (seed >> 16) as i16;
915 fin[i].im = (val as f32) / 256.0;
916 }
917 rdft.do_rdft(&fin, &mut fout1);
9e78289c 918 let mut irdft = RDFTBuilder::new_rdft(fin.len(), false, true);
c3528afd
KS
919 irdft.do_rdft_inplace(&mut fout1);
920
921 for i in 0..fin.len() {
922 let tst = fout1[i].scale(0.5/(fout1.len() as f32));
923 assert!((tst.re - fin[i].re).abs() < 1.0);
924 assert!((tst.im - fin[i].im).abs() < 1.0);
925 }
926 }
3f3b5b23 927}