rustfft/algorithm/
mixed_radix.rs

1use std::cmp::max;
2use std::sync::Arc;
3
4use num_complex::Complex;
5use num_traits::Zero;
6use transpose;
7
8use crate::array_utils;
9use crate::{common::FftNum, twiddles, FftDirection};
10use crate::{Direction, Fft, Length};
11
12/// Implementation of the Mixed-Radix FFT algorithm
13///
14/// This algorithm factors a size n FFT into n1 * n2, computes several inner FFTs of size n1 and n2, then combines the
15/// results to get the final answer
16///
17/// ~~~
18/// // Computes a forward FFT of size 1200, using the Mixed-Radix Algorithm
19/// use rustfft::algorithm::MixedRadix;
20/// use rustfft::{Fft, FftPlanner};
21/// use rustfft::num_complex::Complex;
22///
23/// let mut buffer = vec![Complex{ re: 0.0f32, im: 0.0f32 }; 1200];
24///
25/// // we need to find an n1 and n2 such that n1 * n2 == 1200
26/// // n1 = 30 and n2 = 40 satisfies this
27/// let mut planner = FftPlanner::new();
28/// let inner_fft_n1 = planner.plan_fft_forward(30);
29/// let inner_fft_n2 = planner.plan_fft_forward(40);
30///
31/// // the mixed radix FFT length will be inner_fft_n1.len() * inner_fft_n2.len() = 1200
32/// let fft = MixedRadix::new(inner_fft_n1, inner_fft_n2);
33/// fft.process(&mut buffer);
34/// ~~~
35pub struct MixedRadix<T> {
36    twiddles: Box<[Complex<T>]>,
37
38    width_size_fft: Arc<dyn Fft<T>>,
39    width: usize,
40
41    height_size_fft: Arc<dyn Fft<T>>,
42    height: usize,
43
44    inplace_scratch_len: usize,
45    outofplace_scratch_len: usize,
46    immut_scratch_len: usize,
47
48    direction: FftDirection,
49}
50
51impl<T: FftNum> MixedRadix<T> {
52    /// Creates a FFT instance which will process inputs/outputs of size `width_fft.len() * height_fft.len()`
53    pub fn new(width_fft: Arc<dyn Fft<T>>, height_fft: Arc<dyn Fft<T>>) -> Self {
54        assert_eq!(
55            width_fft.fft_direction(), height_fft.fft_direction(),
56            "width_fft and height_fft must have the same direction. got width direction={}, height direction={}",
57            width_fft.fft_direction(), height_fft.fft_direction());
58
59        let direction = width_fft.fft_direction();
60
61        let width = width_fft.len();
62        let height = height_fft.len();
63
64        let len = width * height;
65
66        let mut twiddles = vec![Complex::zero(); len];
67        for (x, twiddle_chunk) in twiddles.chunks_exact_mut(height).enumerate() {
68            for (y, twiddle_element) in twiddle_chunk.iter_mut().enumerate() {
69                *twiddle_element = twiddles::compute_twiddle(x * y, len, direction);
70            }
71        }
72
73        // Collect some data about what kind of scratch space our inner FFTs need
74        let height_inplace_scratch = height_fft.get_inplace_scratch_len();
75        let width_inplace_scratch = width_fft.get_inplace_scratch_len();
76        let width_outofplace_scratch = width_fft.get_outofplace_scratch_len();
77
78        // Computing the scratch we'll require is a somewhat confusing process.
79        // When we compute an out-of-place FFT, both of our inner FFTs are in-place
80        // When we compute an inplace FFT, our inner width FFT will be inplace, and our height FFT will be out-of-place
81        // For the out-of-place FFT, one of 2 things can happen regarding scratch:
82        //      - If the required scratch of both FFTs is <= self.len(), then we can use the input or output buffer as scratch, and so we need 0 extra scratch
83        //      - If either of the inner FFTs require more, then we'll have to request an entire scratch buffer for the inner FFTs,
84        //          whose size is the max of the two inner FFTs' required scratch
85        let max_inner_inplace_scratch = max(height_inplace_scratch, width_inplace_scratch);
86        let outofplace_scratch_len = if max_inner_inplace_scratch > len {
87            max_inner_inplace_scratch
88        } else {
89            0
90        };
91
92        // For the in-place FFT, again the best case is that we can just bounce data around between internal buffers, and the only inplace scratch we need is self.len()
93        // If our width fft's OOP FFT requires any scratch, then we can tack that on the end of our own scratch, and use split_at_mut to separate our own from our internal FFT's
94        // Likewise, if our height inplace FFT requires more inplace scracth than self.len(), we can tack that on to the end of our own inplace scratch.
95        // Thus, the total inplace scratch is our own length plus the max of what the two inner FFTs will need
96        let inplace_scratch_len = len
97            + max(
98                if height_inplace_scratch > len {
99                    height_inplace_scratch
100                } else {
101                    0
102                },
103                width_outofplace_scratch,
104            );
105
106        let immut_scratch_len = max(
107            len + width_fft.get_inplace_scratch_len(),
108            height_fft.get_inplace_scratch_len(),
109        );
110
111        Self {
112            twiddles: twiddles.into_boxed_slice(),
113
114            width_size_fft: width_fft,
115            width: width,
116
117            height_size_fft: height_fft,
118            height: height,
119
120            inplace_scratch_len,
121            outofplace_scratch_len,
122            immut_scratch_len,
123
124            direction,
125        }
126    }
127
128    fn perform_fft_inplace(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
129        // SIX STEP FFT:
130        let (scratch, inner_scratch) = scratch.split_at_mut(self.len());
131
132        // STEP 1: transpose
133        transpose::transpose(buffer, scratch, self.width, self.height);
134
135        // STEP 2: perform FFTs of size `height`
136        let height_scratch = if inner_scratch.len() > buffer.len() {
137            &mut inner_scratch[..]
138        } else {
139            &mut buffer[..]
140        };
141        self.height_size_fft
142            .process_with_scratch(scratch, height_scratch);
143
144        // STEP 3: Apply twiddle factors
145        for (element, twiddle) in scratch.iter_mut().zip(self.twiddles.iter()) {
146            *element = *element * twiddle;
147        }
148
149        // STEP 4: transpose again
150        transpose::transpose(scratch, buffer, self.height, self.width);
151
152        // STEP 5: perform FFTs of size `width`
153        self.width_size_fft
154            .process_outofplace_with_scratch(buffer, scratch, inner_scratch);
155
156        // STEP 6: transpose again
157        transpose::transpose(scratch, buffer, self.width, self.height);
158    }
159
160    fn perform_fft_immut(
161        &self,
162        input: &[Complex<T>],
163        output: &mut [Complex<T>],
164        scratch_raw: &mut [Complex<T>],
165    ) {
166        // STEP 1: transpose
167        transpose::transpose(input, output, self.width, self.height);
168
169        // STEP 2: perform FFTs of size `height`
170        self.height_size_fft
171            .process_with_scratch(output, scratch_raw);
172
173        // STEP 3: Apply twiddle factors
174        for (element, twiddle) in output.iter_mut().zip(self.twiddles.iter()) {
175            *element = *element * twiddle;
176        }
177
178        let (scratch, inner_scratch) = scratch_raw.split_at_mut(self.len());
179
180        // STEP 4: transpose again
181        transpose::transpose(output, scratch, self.height, self.width);
182
183        // STEP 5: perform FFTs of size `width`
184        self.width_size_fft
185            .process_with_scratch(scratch, inner_scratch);
186
187        // STEP 6: transpose again
188        transpose::transpose(scratch, output, self.width, self.height);
189    }
190
191    fn perform_fft_out_of_place(
192        &self,
193        input: &mut [Complex<T>],
194        output: &mut [Complex<T>],
195        scratch: &mut [Complex<T>],
196    ) {
197        // SIX STEP FFT:
198
199        // STEP 1: transpose
200        transpose::transpose(input, output, self.width, self.height);
201
202        // STEP 2: perform FFTs of size `height`
203        let height_scratch = if scratch.len() > input.len() {
204            &mut scratch[..]
205        } else {
206            &mut input[..]
207        };
208        self.height_size_fft
209            .process_with_scratch(output, height_scratch);
210
211        // STEP 3: Apply twiddle factors
212        for (element, twiddle) in output.iter_mut().zip(self.twiddles.iter()) {
213            *element = *element * twiddle;
214        }
215
216        // STEP 4: transpose again
217        transpose::transpose(output, input, self.height, self.width);
218
219        // STEP 5: perform FFTs of size `width`
220        let width_scratch = if scratch.len() > output.len() {
221            &mut scratch[..]
222        } else {
223            &mut output[..]
224        };
225        self.width_size_fft
226            .process_with_scratch(input, width_scratch);
227
228        // STEP 6: transpose again
229        transpose::transpose(input, output, self.width, self.height);
230    }
231}
232boilerplate_fft!(
233    MixedRadix,
234    |this: &MixedRadix<_>| this.twiddles.len(),
235    |this: &MixedRadix<_>| this.inplace_scratch_len,
236    |this: &MixedRadix<_>| this.outofplace_scratch_len,
237    |this: &MixedRadix<_>| this.immut_scratch_len
238);
239
240/// Implementation of the Mixed-Radix FFT algorithm, specialized for smaller input sizes
241///
242/// This algorithm factors a size n FFT into n1 * n2, computes several inner FFTs of size n1 and n2, then combines the
243/// results to get the final answer
244///
245/// ~~~
246/// // Computes a forward FFT of size 40 using MixedRadixSmall
247/// use std::sync::Arc;
248/// use rustfft::algorithm::MixedRadixSmall;
249/// use rustfft::algorithm::butterflies::{Butterfly5, Butterfly8};
250/// use rustfft::{Fft, FftDirection};
251/// use rustfft::num_complex::Complex;
252///
253/// let len = 40;
254///
255/// let mut buffer = vec![Complex{ re: 0.0f32, im: 0.0f32 }; len];
256///
257/// // we need to find an n1 and n2 such that n1 * n2 == 40
258/// // n1 = 5 and n2 = 8 satisfies this
259/// let inner_fft_n1 = Arc::new(Butterfly5::new(FftDirection::Forward));
260/// let inner_fft_n2 = Arc::new(Butterfly8::new(FftDirection::Forward));
261///
262/// // the mixed radix FFT length will be inner_fft_n1.len() * inner_fft_n2.len() = 40
263/// let fft = MixedRadixSmall::new(inner_fft_n1, inner_fft_n2);
264/// fft.process(&mut buffer);
265/// ~~~
266pub struct MixedRadixSmall<T> {
267    twiddles: Box<[Complex<T>]>,
268
269    width_size_fft: Arc<dyn Fft<T>>,
270    width: usize,
271
272    height_size_fft: Arc<dyn Fft<T>>,
273    height: usize,
274
275    direction: FftDirection,
276}
277
278impl<T: FftNum> MixedRadixSmall<T> {
279    /// Creates a FFT instance which will process inputs/outputs of size `width_fft.len() * height_fft.len()`
280    pub fn new(width_fft: Arc<dyn Fft<T>>, height_fft: Arc<dyn Fft<T>>) -> Self {
281        assert_eq!(
282            width_fft.fft_direction(), height_fft.fft_direction(),
283            "width_fft and height_fft must have the same direction. got width direction={}, height direction={}",
284            width_fft.fft_direction(), height_fft.fft_direction());
285
286        // Verify that the inner FFTs don't require out-of-place scratch, and only arequire a small amount of inplace scratch
287        let width = width_fft.len();
288        let height = height_fft.len();
289        let len = width * height;
290
291        assert_eq!(width_fft.get_outofplace_scratch_len(), 0, "MixedRadixSmall should only be used with algorithms that require 0 out-of-place scratch. Width FFT (len={}) requires {}, should require 0", width, width_fft.get_outofplace_scratch_len());
292        assert_eq!(height_fft.get_outofplace_scratch_len(), 0, "MixedRadixSmall should only be used with algorithms that require 0 out-of-place scratch. Height FFT (len={}) requires {}, should require 0", height, height_fft.get_outofplace_scratch_len());
293
294        assert!(width_fft.get_inplace_scratch_len() <= width, "MixedRadixSmall should only be used with algorithms that require little inplace scratch. Width FFT (len={}) requires {}, should require {} or less", width, width_fft.get_inplace_scratch_len(), width);
295        assert!(height_fft.get_inplace_scratch_len() <= height, "MixedRadixSmall should only be used with algorithms that require little inplace scratch. Height FFT (len={}) requires {}, should require {} or less", height, height_fft.get_inplace_scratch_len(), height);
296
297        let direction = width_fft.fft_direction();
298
299        let mut twiddles = vec![Complex::zero(); len];
300        for (x, twiddle_chunk) in twiddles.chunks_exact_mut(height).enumerate() {
301            for (y, twiddle_element) in twiddle_chunk.iter_mut().enumerate() {
302                *twiddle_element = twiddles::compute_twiddle(x * y, len, direction);
303            }
304        }
305
306        Self {
307            twiddles: twiddles.into_boxed_slice(),
308
309            width_size_fft: width_fft,
310            width: width,
311
312            height_size_fft: height_fft,
313            height: height,
314
315            direction,
316        }
317    }
318
319    fn perform_fft_inplace(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
320        // SIX STEP FFT:
321        // STEP 1: transpose
322        unsafe { array_utils::transpose_small(self.width, self.height, buffer, scratch) };
323
324        // STEP 2: perform FFTs of size `height`
325        self.height_size_fft.process_with_scratch(scratch, buffer);
326
327        // STEP 3: Apply twiddle factors
328        for (element, twiddle) in scratch.iter_mut().zip(self.twiddles.iter()) {
329            *element = *element * twiddle;
330        }
331
332        // STEP 4: transpose again
333        unsafe { array_utils::transpose_small(self.height, self.width, scratch, buffer) };
334
335        // STEP 5: perform FFTs of size `width`
336        self.width_size_fft
337            .process_outofplace_with_scratch(buffer, scratch, &mut []);
338
339        // STEP 6: transpose again
340        unsafe { array_utils::transpose_small(self.width, self.height, scratch, buffer) };
341    }
342
343    fn perform_fft_immut(
344        &self,
345        input: &[Complex<T>],
346        output: &mut [Complex<T>],
347        scratch: &mut [Complex<T>],
348    ) {
349        // SIX STEP FFT:
350        // STEP 1: transpose
351        unsafe { array_utils::transpose_small(self.width, self.height, input, output) };
352
353        // STEP 2: perform FFTs of size `height`
354        self.height_size_fft.process_with_scratch(output, scratch);
355
356        // STEP 3: Apply twiddle factors
357        for (element, twiddle) in output.iter_mut().zip(self.twiddles.iter()) {
358            *element = *element * twiddle;
359        }
360
361        // STEP 4: transpose again
362        unsafe { array_utils::transpose_small(self.height, self.width, output, scratch) };
363
364        // STEP 5: perform FFTs of size `width`
365        self.width_size_fft.process_with_scratch(scratch, output);
366
367        // STEP 6: transpose again
368        unsafe { array_utils::transpose_small(self.width, self.height, scratch, output) };
369    }
370
371    fn perform_fft_out_of_place(
372        &self,
373        input: &mut [Complex<T>],
374        output: &mut [Complex<T>],
375        _scratch: &mut [Complex<T>],
376    ) {
377        // SIX STEP FFT:
378        // STEP 1: transpose
379        unsafe { array_utils::transpose_small(self.width, self.height, input, output) };
380
381        // STEP 2: perform FFTs of size `height`
382        self.height_size_fft.process_with_scratch(output, input);
383
384        // STEP 3: Apply twiddle factors
385        for (element, twiddle) in output.iter_mut().zip(self.twiddles.iter()) {
386            *element = *element * twiddle;
387        }
388
389        // STEP 4: transpose again
390        unsafe { array_utils::transpose_small(self.height, self.width, output, input) };
391
392        // STEP 5: perform FFTs of size `width`
393        self.width_size_fft.process_with_scratch(input, output);
394
395        // STEP 6: transpose again
396        unsafe { array_utils::transpose_small(self.width, self.height, input, output) };
397    }
398}
399boilerplate_fft!(
400    MixedRadixSmall,
401    |this: &MixedRadixSmall<_>| this.twiddles.len(),
402    |this: &MixedRadixSmall<_>| this.len(),
403    |_| 0,
404    |this: &MixedRadixSmall<_>| this.len()
405);
406
407#[cfg(test)]
408mod unit_tests {
409    use super::*;
410    use crate::test_utils::check_fft_algorithm;
411    use crate::{algorithm::Dft, test_utils::BigScratchAlgorithm};
412    use num_traits::Zero;
413    use std::sync::Arc;
414
415    #[test]
416    fn test_mixed_radix() {
417        for width in 1..7 {
418            for height in 1..7 {
419                test_mixed_radix_with_lengths(width, height, FftDirection::Forward);
420                test_mixed_radix_with_lengths(width, height, FftDirection::Inverse);
421            }
422        }
423    }
424
425    #[test]
426    fn test_mixed_radix_small() {
427        for width in 2..7 {
428            for height in 2..7 {
429                test_mixed_radix_small_with_lengths(width, height, FftDirection::Forward);
430                test_mixed_radix_small_with_lengths(width, height, FftDirection::Inverse);
431            }
432        }
433    }
434
435    fn test_mixed_radix_with_lengths(width: usize, height: usize, direction: FftDirection) {
436        let width_fft = Arc::new(Dft::new(width, direction)) as Arc<dyn Fft<f32>>;
437        let height_fft = Arc::new(Dft::new(height, direction)) as Arc<dyn Fft<f32>>;
438
439        let fft = MixedRadix::new(width_fft, height_fft);
440
441        check_fft_algorithm(&fft, width * height, direction);
442    }
443
444    fn test_mixed_radix_small_with_lengths(width: usize, height: usize, direction: FftDirection) {
445        let width_fft = Arc::new(Dft::new(width, direction)) as Arc<dyn Fft<f32>>;
446        let height_fft = Arc::new(Dft::new(height, direction)) as Arc<dyn Fft<f32>>;
447
448        let fft = MixedRadixSmall::new(width_fft, height_fft);
449
450        check_fft_algorithm(&fft, width * height, direction);
451    }
452
453    // Verify that the mixed radix algorithm correctly provides scratch space to inner FFTs
454    #[test]
455    fn test_mixed_radix_inner_scratch() {
456        let scratch_lengths = [1, 5, 25];
457
458        let mut inner_ffts = Vec::new();
459
460        for &len in &scratch_lengths {
461            for &inplace_scratch in &scratch_lengths {
462                for &outofplace_scratch in &scratch_lengths {
463                    for &immut_scratch in &scratch_lengths {
464                        inner_ffts.push(Arc::new(BigScratchAlgorithm {
465                            len,
466                            inplace_scratch,
467                            outofplace_scratch,
468                            immut_scratch,
469                            direction: FftDirection::Forward,
470                        }) as Arc<dyn Fft<f32>>);
471                    }
472                }
473            }
474        }
475
476        for width_fft in inner_ffts.iter() {
477            for height_fft in inner_ffts.iter() {
478                let fft = MixedRadix::new(Arc::clone(width_fft), Arc::clone(height_fft));
479
480                let mut inplace_buffer = vec![Complex::zero(); fft.len()];
481                let mut inplace_scratch = vec![Complex::zero(); fft.get_inplace_scratch_len()];
482
483                fft.process_with_scratch(&mut inplace_buffer, &mut inplace_scratch);
484
485                let mut outofplace_input = vec![Complex::zero(); fft.len()];
486                let mut outofplace_output = vec![Complex::zero(); fft.len()];
487                let mut outofplace_scratch =
488                    vec![Complex::zero(); fft.get_outofplace_scratch_len()];
489                fft.process_outofplace_with_scratch(
490                    &mut outofplace_input,
491                    &mut outofplace_output,
492                    &mut outofplace_scratch,
493                );
494
495                let immut_input = vec![Complex::zero(); fft.len()];
496                let mut immut_output = vec![Complex::zero(); fft.len()];
497                let mut immut_scratch = vec![Complex::zero(); fft.get_immutable_scratch_len()];
498
499                fft.process_immutable_with_scratch(
500                    &immut_input,
501                    &mut immut_output,
502                    &mut immut_scratch,
503                );
504            }
505        }
506    }
507}
rustfft/algorithm/mixed_radix.rs

rustfft/algorithm/
mixed_radix.rs
