]>
Commit | Line | Data |
---|---|---|
4e034a32 | 1 | //! FFT and RDFT implementation. |
3f3b5b23 | 2 | use std::f32::{self, consts}; |
03d914bb | 3 | use std::ops::{Not, Neg, Add, AddAssign, Sub, SubAssign, Mul, MulAssign, Div}; |
3f3b5b23 KS |
4 | use std::fmt; |
5 | ||
4e034a32 | 6 | /// Complex number. |
b01971f2 | 7 | #[repr(C)] |
3f3b5b23 KS |
8 | #[derive(Debug,Clone,Copy,PartialEq)] |
9 | pub 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 | ||
16 | impl 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 | ||
39 | impl Neg for FFTComplex { | |
40 | type Output = FFTComplex; | |
41 | fn neg(self) -> Self::Output { | |
42 | FFTComplex { re: -self.re, im: -self.im } | |
43 | } | |
44 | } | |
45 | ||
46 | impl Not for FFTComplex { | |
47 | type Output = FFTComplex; | |
48 | fn not(self) -> Self::Output { | |
49 | FFTComplex { re: self.re, im: -self.im } | |
50 | } | |
51 | } | |
52 | ||
53 | impl 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 | ||
60 | impl AddAssign for FFTComplex { | |
61 | fn add_assign(&mut self, other: Self) { | |
62 | self.re += other.re; | |
63 | self.im += other.im; | |
64 | } | |
65 | } | |
66 | ||
67 | impl 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 | ||
74 | impl SubAssign for FFTComplex { | |
75 | fn sub_assign(&mut self, other: Self) { | |
76 | self.re -= other.re; | |
77 | self.im -= other.im; | |
78 | } | |
79 | } | |
80 | ||
81 | impl 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 | ||
89 | impl 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 |
98 | impl Div for FFTComplex { |
99 | type Output = FFTComplex; | |
100 | fn div(self, other: Self) -> Self::Output { | |
101 | self * other.reciprocal() | |
102 | } | |
103 | } | |
104 | ||
3f3b5b23 KS |
105 | impl 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 |
112 | pub 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 |
115 | pub 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 |
133 | struct 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 |
142 | struct FFTGeneric {} |
143 | ||
144 | const FFT3_CONST: f32 = 0.86602540378443864677; | |
145 | const FFT5_CONST1: FFTComplex = FFTComplex { re: 0.80901699437494742410, im: 0.58778525229247312915 }; | |
146 | const FFT5_CONST2: FFTComplex = FFTComplex { re: 0.30901699437494742411, im: 0.95105651629515357211 }; | |
147 | ||
148 | fn 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 | ||
157 | impl 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 | ||
265 | struct FFTSplitRadix {} | |
266 | ||
267 | impl 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 | ||
386 | struct FFT15 {} | |
3f3b5b23 | 387 | |
9e78289c KS |
388 | const FFT15_INSWAP: [usize; 20] = [ 0, 5, 10, 42, 3, 8, 13, 42, 6, 11, 1, 42, 9, 14, 4, 42, 12, 2, 7, 42 ]; |
389 | const 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 | ||
391 | impl 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 | ||
486 | enum FFTMode { | |
487 | Generic(usize), | |
488 | SplitRadix(u8), | |
489 | Prime15, | |
490 | } | |
491 | ||
492 | impl 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 |
551 | pub struct FFT { |
552 | perms: Vec<usize>, | |
553 | swaps: Vec<usize>, | |
554 | ffts: Vec<(FFTMode, FFTData)>, | |
555 | } | |
556 | ||
557 | impl 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 |
641 | pub 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 | ||
659 | fn 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 | |
664 | fn 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 | 683 | fn 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 | ||
706 | impl 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 |
762 | pub 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 | 770 | fn crossadd(a: FFTComplex, b: FFTComplex) -> FFTComplex { |
c3528afd KS |
771 | FFTComplex { re: a.re + b.re, im: a.im - b.im } |
772 | } | |
773 | ||
774 | impl 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 |
832 | pub struct RDFTBuilder { |
833 | } | |
834 | ||
835 | impl 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)] | |
850 | mod 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 | } |