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