feanor_math/algorithms/fft/
radix3.rs

1use std::alloc::{Allocator, Global};
2
3use crate::algorithms::fft::FFTAlgorithm;
4use crate::rings::float_complex::Complex64Base;
5use crate::algorithms::fft::complex_fft::*;
6use crate::algorithms::unity_root::{get_prim_root_of_unity_zn, is_prim_root_of_unity};
7use crate::divisibility::{DivisibilityRing, DivisibilityRingStore};
8use crate::homomorphism::*;
9use crate::primitive_int::StaticRing;
10use crate::seq::SwappableVectorViewMut;
11use crate::rings::zn::*;
12use crate::ring::*;
13use crate::seq::VectorFn;
14
15///
16/// Implementation of the Cooley-Tukey FFT algorithm for power-of-three lengths.
17/// 
18#[stability::unstable(feature = "enable")]
19pub struct CooleyTukeyRadix3FFT<R_main, R_twiddle, H, A = Global>
20    where R_main: ?Sized + RingBase,
21        R_twiddle: ?Sized + RingBase + DivisibilityRing,
22        H: Homomorphism<R_twiddle, R_main>,
23        A: Allocator
24{
25    log3_n: usize,
26    hom: H,
27    twiddles: Vec<Vec<R_twiddle::Element>>,
28    inv_twiddles: Vec<Vec<R_twiddle::Element>>,
29    third_root_of_unity: R_twiddle::Element,
30    main_root_of_unity: R_main::Element,
31    allocator: A
32}
33
34const ZZ: StaticRing<i64> = StaticRing::RING;
35
36#[inline(never)]
37fn butterfly_loop<T, S, F>(log3_n: usize, data: &mut [T], step: usize, twiddles: &[S], butterfly: F)
38    where F: Fn(&mut T, &mut T, &mut T, &S, &S) + Clone
39{
40    assert_eq!(ZZ.pow(3, log3_n) as usize, data.len());
41    assert!(step < log3_n);
42    assert_eq!(2 * ZZ.pow(3, step) as usize, twiddles.len());
43    
44    let stride = ZZ.pow(3, log3_n - step - 1) as usize;
45    assert!(data.len() % (3 * stride) == 0);
46    assert_eq!(twiddles.as_chunks::<2>().0.len(), data.chunks_mut(3 * stride).len());
47
48    if stride == 1 {
49        for ([twiddle1, twiddle2], butterfly_data) in twiddles.as_chunks::<2>().0.iter().zip(data.as_chunks_mut::<3>().0.iter_mut()) {
50            let [a, b, c] = butterfly_data.each_mut();
51            butterfly(a, b, c, twiddle1, twiddle2);
52        }
53    } else {
54        for ([twiddle1, twiddle2], butterfly_data) in twiddles.as_chunks::<2>().0.iter().zip(data.chunks_mut(3 * stride)) {
55            let (first, rest) = butterfly_data.split_at_mut(stride);
56            let (second, third) = rest.split_at_mut(stride);
57            for ((a, b), c) in first.iter_mut().zip(second.iter_mut()).zip(third.iter_mut()) {
58                butterfly(a, b, c, twiddle1, twiddle2);
59            }
60        }
61    }
62}
63
64fn threeadic_reverse(mut number: usize, log3_n: usize) -> usize {
65    debug_assert!((number as i64) < ZZ.pow(3, log3_n));
66    let mut result = 0;
67    for _ in 0..log3_n {
68        let (quo, rem) = (number / 3, number % 3);
69        result = 3 * result + rem;
70        number = quo;
71    }
72    assert_eq!(0, number);
73    return result;
74}
75
76impl<R> CooleyTukeyRadix3FFT<R::Type, R::Type, Identity<R>>
77    where R: RingStore,
78        R::Type: DivisibilityRing
79{
80    ///
81    /// Creates an [`CooleyTukeyRadix3FFT`] for a prime field, assuming it has a characteristic
82    /// congruent to 1 modulo `3^lo32_n`.
83    /// 
84    #[stability::unstable(feature = "enable")]
85    pub fn for_zn(ring: R, log3_n: usize) -> Option<Self>
86        where R::Type: ZnRing
87    {
88        let n = ZZ.pow(3, log3_n);
89        let root_of_unity = get_prim_root_of_unity_zn(&ring, n as usize)?;
90        return Some(Self::new_with_hom(ring.into_identity(), root_of_unity, log3_n));
91    }
92}
93
94impl<R_main, R_twiddle, H> CooleyTukeyRadix3FFT<R_main, R_twiddle, H>
95    where R_main: ?Sized + RingBase,
96        R_twiddle: ?Sized + RingBase + DivisibilityRing,
97        H: Homomorphism<R_twiddle, R_main>
98{
99    ///
100    /// Creates an [`CooleyTukeyRadix3FFT`] for the given rings, using the given root of unity.
101    /// 
102    /// Instead of a ring, this function takes a homomorphism `R -> S`. Twiddle factors that are
103    /// precomputed will be stored as elements of `R`, while the main FFT computations will be 
104    /// performed in `S`. This allows both implicit ring conversions, and using patterns like 
105    /// [`zn_64::ZnFastmul`] to precompute some data for better performance.
106    /// 
107    /// Do not use this for approximate rings, as computing the powers of `root_of_unity`
108    /// will incur avoidable precision loss.
109    /// 
110    #[stability::unstable(feature = "enable")]
111    pub fn new_with_hom(hom: H, zeta: R_twiddle::Element, log3_n: usize) -> Self {
112        let ring = hom.domain();
113        let pow_zeta = |i: i64| if i < 0 {
114            ring.invert(&ring.pow(ring.clone_el(&zeta), (-i).try_into().unwrap())).unwrap()
115        } else {
116            ring.pow(ring.clone_el(&zeta), i.try_into().unwrap())
117        };
118        let result = CooleyTukeyRadix3FFT::create(&hom, pow_zeta, log3_n, Global);
119        return Self {
120            allocator: result.allocator,
121            inv_twiddles: result.inv_twiddles,
122            log3_n: result.log3_n,
123            main_root_of_unity: result.main_root_of_unity,
124            third_root_of_unity: result.third_root_of_unity,
125            twiddles: result.twiddles,
126            hom: hom
127        };
128    }
129}
130
131impl<R_main, R_twiddle, H, A> CooleyTukeyRadix3FFT<R_main, R_twiddle, H, A>
132    where R_main: ?Sized + RingBase,
133        R_twiddle: ?Sized + RingBase + DivisibilityRing,
134        H: Homomorphism<R_twiddle, R_main>,
135        A: Allocator
136{
137    ///
138    /// Most general way to create a [`CooleyTukeyRadix3FFT`].
139    /// 
140    /// The given closure should, on input `i`, return `z^i` for a primitive
141    /// `3^log3_n`-th root of unity. The given allocator is used to copy the input
142    /// data in cases where the input data layout is not optimal for the algorithm
143    /// 
144    #[stability::unstable(feature = "enable")]
145    pub fn create<F>(hom: H, mut root_of_unity_pow: F, log3_n: usize, allocator: A) -> Self 
146        where F: FnMut(i64) -> R_twiddle::Element
147    {
148        let n = ZZ.pow(3, log3_n) as usize;
149        assert!(hom.codomain().get_ring().is_approximate() || is_prim_root_of_unity(hom.codomain(), &hom.map(root_of_unity_pow(1)), n));
150        
151        return Self {
152            main_root_of_unity: hom.map(root_of_unity_pow(1)),
153            log3_n: log3_n, 
154            twiddles: Self::create_twiddle_list(hom.domain(), log3_n, &mut root_of_unity_pow), 
155            inv_twiddles: Self::create_inv_twiddle_list(hom.domain(), log3_n, &mut root_of_unity_pow),
156            third_root_of_unity: root_of_unity_pow(2 * n as i64 / 3),
157            hom: hom,
158            allocator: allocator
159        };
160    }
161    
162    ///
163    /// Replaces the ring that this object can compute FFTs over, assuming that the current
164    /// twiddle factors can be mapped into the new ring with the given homomorphism.
165    /// 
166    /// In particular, this function does not recompute twiddles, but uses a different
167    /// homomorphism to map the current twiddles into a new ring. Hence, it is extremely
168    /// cheap. 
169    /// 
170    #[stability::unstable(feature = "enable")]
171    pub fn change_ring<R_new: ?Sized + RingBase, H_new: Homomorphism<R_twiddle, R_new>>(self, new_hom: H_new) -> (CooleyTukeyRadix3FFT<R_new, R_twiddle, H_new, A>, H) {
172        let ring = new_hom.codomain();
173        let root_of_unity = if self.log3_n == 0 {
174            new_hom.codomain().one()
175        } else if self.log3_n == 1 {
176            let root_of_unity = self.hom.domain().pow(self.hom.domain().clone_el(&self.third_root_of_unity), 2);
177            debug_assert!(self.ring().eq_el(&self.hom.map_ref(&root_of_unity), self.root_of_unity(self.hom.codomain())));
178            new_hom.map(root_of_unity)
179        } else {
180            let root_of_unity = &self.inv_twiddles[self.log3_n - 1][threeadic_reverse(1, self.log3_n - 1)];
181            debug_assert!(self.ring().eq_el(&self.hom.map_ref(root_of_unity), self.root_of_unity(self.hom.codomain())));
182            new_hom.map_ref(root_of_unity)
183        };
184        assert!(ring.is_commutative());
185        assert!(ring.get_ring().is_approximate() || is_prim_root_of_unity(&ring, &root_of_unity, self.len()));
186
187        return (
188            CooleyTukeyRadix3FFT {
189                twiddles: self.twiddles,
190                inv_twiddles: self.inv_twiddles,
191                main_root_of_unity: root_of_unity, 
192                third_root_of_unity: self.third_root_of_unity,
193                hom: new_hom, 
194                log3_n: self.log3_n, 
195                allocator: self.allocator
196            },
197            self.hom
198        );
199    }
200
201    ///
202    /// Returns the ring over which this object can compute FFTs.
203    /// 
204    #[stability::unstable(feature = "enable")]
205    pub fn ring<'a>(&'a self) -> &'a <H as Homomorphism<R_twiddle, R_main>>::CodomainStore {
206        self.hom.codomain()
207    }
208
209    ///
210    /// Returns a reference to the allocator currently used for temporary allocations by this FFT.
211    /// 
212    #[stability::unstable(feature = "enable")]
213    pub fn allocator(&self) -> &A {
214        &self.allocator
215    }
216
217    ///
218    /// Replaces the allocator used for temporary allocations by this FFT.
219    /// 
220    #[stability::unstable(feature = "enable")]
221    pub fn with_allocator<A_new: Allocator>(self, allocator: A_new) -> CooleyTukeyRadix3FFT<R_main, R_twiddle, H, A_new> {
222        CooleyTukeyRadix3FFT {
223            twiddles: self.twiddles,
224            inv_twiddles: self.inv_twiddles,
225            main_root_of_unity: self.main_root_of_unity, 
226            third_root_of_unity: self.third_root_of_unity,
227            hom: self.hom, 
228            log3_n: self.log3_n, 
229            allocator: allocator
230        }
231    }
232
233    ///
234    /// Returns a reference to the homomorphism that is used to map the stored twiddle
235    /// factors into main ring, over which FFTs are computed.
236    /// 
237    #[stability::unstable(feature = "enable")]
238    pub fn hom<'a>(&'a self) -> &'a H {
239        &self.hom
240    }
241
242    fn create_twiddle_list<F>(ring: &H::DomainStore, log3_n: usize, mut pow_zeta: F) -> Vec<Vec<R_twiddle::Element>>
243        where F: FnMut(i64) -> R_twiddle::Element
244    {
245        let n = ZZ.pow(3, log3_n);
246        let third_root_of_unity = pow_zeta(-(n / 3));
247        let mut result: Vec<_> = (0..log3_n).map(|_| Vec::new()).collect();
248        for i in 0..log3_n {
249            let current = &mut result[i];
250            for j in 0..ZZ.pow(3, i) {
251                let base_twiddle = pow_zeta(-(threeadic_reverse(j as usize, log3_n - 1) as i64));
252                current.push(ring.clone_el(&base_twiddle));
253                current.push(ring.pow(ring.mul_ref_snd(base_twiddle, &third_root_of_unity), 2));
254            }
255        }
256        return result;
257    }
258    
259    fn create_inv_twiddle_list<F>(ring: &H::DomainStore, log3_n: usize, mut pow_zeta: F) -> Vec<Vec<R_twiddle::Element>>
260        where F: FnMut(i64) -> R_twiddle::Element
261    {
262        let mut result: Vec<_> = (0..log3_n).map(|_| Vec::new()).collect();
263        for i in 0..log3_n {
264            let current = &mut result[i];
265            for j in 0..ZZ.pow(3, i) {
266                let base_twiddle = pow_zeta(threeadic_reverse(j as usize, log3_n - 1) as i64);
267                current.push(ring.clone_el(&base_twiddle));
268                current.push(ring.pow(base_twiddle, 2));
269            }
270        }
271        return result;
272    }
273
274    fn butterfly_step_main<const INV: bool>(&self, data: &mut [R_main::Element], step: usize) {
275        let twiddles = if INV {
276            &self.inv_twiddles[step]
277        } else {
278            &self.twiddles[step]
279        };
280        let third_root_of_unity = &self.third_root_of_unity;
281        // let start = std::time::Instant::now();
282        butterfly_loop(self.log3_n, data, step, twiddles, |x, y, z, twiddle1, twiddle2| if INV {
283            <R_main as CooleyTukeyRadix3Butterfly<R_twiddle>>::inv_butterfly(&self.hom, x, y, z, &third_root_of_unity, twiddle1, twiddle2)
284        } else {
285            <R_main as CooleyTukeyRadix3Butterfly<R_twiddle>>::butterfly(&self.hom, x, y, z, &third_root_of_unity, twiddle1, twiddle2)
286        });
287        // let end = std::time::Instant::now();
288        // BUTTERFLY_RADIX3_TIMES[step].fetch_add((end - start).as_micros() as usize, std::sync::atomic::Ordering::Relaxed);
289    }
290
291    fn fft_impl(&self, data: &mut [R_main::Element]) {
292        for i in 0..data.len() {
293            <R_main as CooleyTukeyRadix3Butterfly<R_twiddle>>::prepare_for_fft(self.hom.codomain().get_ring(), &mut data[i]);
294        }
295        for step in 0..self.log3_n {
296            self.butterfly_step_main::<false>(data, step);
297        }
298    }
299
300    fn inv_fft_impl(&self, data: &mut [R_main::Element]) {
301        for i in 0..data.len() {
302            <R_main as CooleyTukeyRadix3Butterfly<R_twiddle>>::prepare_for_inv_fft(self.hom.codomain().get_ring(), &mut data[i]);
303        }
304        for step in (0..self.log3_n).rev() {
305            self.butterfly_step_main::<true>(data, step);
306        }
307        let n_inv = self.hom.domain().invert(&self.hom.domain().int_hom().map(self.len() as i32)).unwrap();
308        for i in 0..data.len() {
309            self.hom.mul_assign_ref_map(&mut data[i], &n_inv);
310        }
311    }
312}
313
314impl<R_main, R_twiddle, H, A> Clone for CooleyTukeyRadix3FFT<R_main, R_twiddle, H, A> 
315    where R_main: ?Sized + RingBase,
316        R_twiddle: ?Sized + RingBase + DivisibilityRing,
317        H: Homomorphism<R_twiddle, R_main> + Clone,
318        A: Allocator + Clone
319{
320    fn clone(&self) -> Self {
321        Self {
322            hom: self.hom.clone(),
323            inv_twiddles: self.inv_twiddles.iter().map(|list| list.iter().map(|x| self.hom.domain().clone_el(x)).collect()).collect(),
324            twiddles: self.twiddles.iter().map(|list| list.iter().map(|x| self.hom.domain().clone_el(x)).collect()).collect(),
325            main_root_of_unity: self.hom.codomain().clone_el(&self.main_root_of_unity),
326            third_root_of_unity: self.hom.domain().clone_el(&self.third_root_of_unity),
327            log3_n: self.log3_n,
328            allocator: self.allocator.clone()
329        }
330    }
331}
332
333impl<R_main, R_twiddle, H, A> FFTAlgorithm<R_main> for CooleyTukeyRadix3FFT<R_main, R_twiddle, H, A>
334    where R_main: ?Sized + RingBase,
335        R_twiddle: ?Sized + RingBase + DivisibilityRing,
336        H: Homomorphism<R_twiddle, R_main>,
337        A: Allocator
338{
339    fn len(&self) -> usize {
340        if self.log3_n == 0 {
341            return 1;
342        }
343        let result = self.twiddles[self.log3_n - 1].len() / 2 * 3;
344        debug_assert_eq!(ZZ.pow(3, self.log3_n) as usize, result);
345        return result;
346    }
347
348    fn unordered_fft<V, S>(&self, mut values: V, ring: S)
349        where V: SwappableVectorViewMut<R_main::Element>,
350            S: RingStore<Type = R_main> + Copy
351    {
352        assert!(ring.get_ring() == self.hom.codomain().get_ring(), "unsupported ring");
353        assert_eq!(self.len(), values.len());
354        if let Some(data) = values.as_slice_mut() {
355            self.fft_impl(data);
356        } else {
357            let mut data = Vec::with_capacity_in(self.len(), &self.allocator);
358            data.extend(values.clone_ring_els(ring).iter());
359            self.fft_impl(&mut data);
360            for (i, x) in data.into_iter().enumerate() {
361                *values.at_mut(i) = x;
362            }
363        }
364    }
365
366    fn unordered_inv_fft<V, S>(&self, mut values: V, ring: S)
367        where V: SwappableVectorViewMut<R_main::Element>,
368            S: RingStore<Type = R_main> + Copy
369    {
370        assert!(ring.get_ring() == self.hom.codomain().get_ring(), "unsupported ring");
371        assert_eq!(self.len(), values.len());
372        if let Some(data) = values.as_slice_mut() {
373            self.inv_fft_impl(data);
374        } else {
375            let mut data = Vec::with_capacity_in(self.len(), &self.allocator);
376            data.extend(values.clone_ring_els(ring).iter());
377            self.inv_fft_impl(&mut data);
378            for (i, x) in data.into_iter().enumerate() {
379                *values.at_mut(i) = x;
380            }
381        }
382    }
383
384    fn root_of_unity<S>(&self, ring: S) -> &R_main::Element
385        where S: RingStore<Type = R_main> + Copy
386    {
387        assert!(ring.get_ring() == self.hom.codomain().get_ring(), "unsupported ring");   
388        &self.main_root_of_unity 
389    }
390
391    fn unordered_fft_permutation(&self, i: usize) -> usize {
392        threeadic_reverse(i, self.log3_n)
393    }
394
395    fn unordered_fft_permutation_inv(&self, i: usize) -> usize {
396        threeadic_reverse(i, self.log3_n)
397    }
398}
399
400#[stability::unstable(feature = "enable")]
401pub trait CooleyTukeyRadix3Butterfly<S: ?Sized + RingBase>: RingBase {
402
403    ///
404    /// Should compute `(a, b, c) := (a + t b + t^2 c, a + t z b + t^2 z^2 c, a + t z^2 b + t^2 z c)`.
405    /// 
406    /// Here `z` is a third root of unity (i.e. `z^2 + z + 1 = 0`) and `t` is the twiddle factor.
407    /// The function should be given `z, t, t^2 z^2`.
408    /// 
409    /// It is guaranteed that the input elements are either outputs of
410    /// [`CooleyTukeyRadix3Butterfly::butterfly()`] or of [`CooleyTukeyRadix3Butterfly::prepare_for_fft()`].
411    /// 
412    fn butterfly<H: Homomorphism<S, Self>>(
413        hom: H, 
414        a: &mut Self::Element, 
415        b: &mut Self::Element, 
416        c: &mut Self::Element, 
417        z: &S::Element,
418        t: &S::Element,
419        t_sqr_z_sqr: &S::Element
420    );
421    
422    ///
423    /// Should compute `(a, b, c) := (a + b + c, t (a + z^2 b + z c), t^2 (a + z b + z^2 c))`.
424    /// 
425    /// Here `z` is a third root of unity (i.e. `z^2 + z + 1 = 0`) and `t` is the twiddle factor.
426    /// The function should be given `z, t, t^2`.
427    /// 
428    /// It is guaranteed that the input elements are either outputs of
429    /// [`CooleyTukeyRadix3Butterfly::inv_butterfly()`] or of [`CooleyTukeyRadix3Butterfly::prepare_for_inv_fft()`].
430    /// 
431    fn inv_butterfly<H: Homomorphism<S, Self>>(
432        hom: H, 
433        a: &mut Self::Element, 
434        b: &mut Self::Element,
435        c: &mut Self::Element,
436        z: &S::Element,
437        t: &S::Element,
438        t_sqr: &S::Element
439    );
440
441    ///
442    /// Possibly pre-processes elements before the FFT starts. Here you can
443    /// bring ring element into a certain form, and assume during [`CooleyTukeyRadix3Butterfly::butterfly()`]
444    /// that the inputs are in this form.
445    /// 
446    #[inline(always)]
447    fn prepare_for_fft(&self, _value: &mut Self::Element) {}
448    
449    ///
450    /// Possibly pre-processes elements before the inverse FFT starts. Here you can
451    /// bring ring element into a certain form, and assume during [`CooleyTukeyRadix3Butterfly::inv_butterfly()`]
452    /// that the inputs are in this form.
453    /// 
454    #[inline(always)]
455    fn prepare_for_inv_fft(&self, _value: &mut Self::Element) {}
456}
457
458impl<R: ?Sized + RingBase, S: ?Sized + RingBase> CooleyTukeyRadix3Butterfly<S> for R {
459
460    default fn butterfly<H: Homomorphism<S, Self>>(
461        hom: H, 
462        a: &mut Self::Element, 
463        b: &mut Self::Element, 
464        c: &mut Self::Element, 
465        z: &S::Element,
466        t: &S::Element,
467        t_sqr_z_sqr: &S::Element
468    ) {
469        let ring = hom.codomain();
470        hom.mul_assign_ref_map(b, t); // this is now `t b`
471        hom.mul_assign_ref_map(c, t_sqr_z_sqr); // this is now `t^2 z^2 c`
472        let b_ = hom.mul_ref_map(b, z); // this is now `t z b`
473        let c_ = hom.mul_ref_map(c, z); // this is now `t^2 c z`
474        let s1 = ring.add_ref(b, &c_); // this is now `t b + t^2 c`
475        let s2 = ring.add_ref(&b_, c); // this is now `t z b + t^2 z^2 c`
476        let s3 = ring.add_ref(&s1, &s2); // this is now `-(t z^2 b + t^2 z c)`
477        *b = ring.add_ref_fst(a, s2); // this is now `a + t z b + t^2 z^2 c`
478        *c = ring.sub_ref_fst(a, s3); // this is now `a + t z^2 b + t^2 z c`
479        ring.add_assign(a, s1); // this is now `a + t b + t^2 c`
480    }
481    
482    default fn inv_butterfly<H: Homomorphism<S, Self>>(
483        hom: H, 
484        a: &mut Self::Element, 
485        b: &mut Self::Element,
486        c: &mut Self::Element,
487        z: &S::Element,
488        t: &S::Element,
489        t_sqr: &S::Element
490    ) {
491        let ring = hom.codomain();
492        let b_ = hom.mul_ref_map(b, z); // this is now `z b`
493        let s1 = ring.add_ref(b, c); // this is now `b + c`
494        let s2 = ring.add_ref(&b_, &c); // this is now `z b + c`
495        let s2_ = hom.mul_ref_snd_map(s2, z); // this is now `z^2 b + z c`
496        let s3 = ring.add_ref(&s1, &s2_); // this is now `-(z b + z^2 c)`
497        *b = ring.add_ref(a, &s2_); // this is now `a + z^2 b + z c`
498        *c = ring.sub_ref(a, &s3); // this is now `a + z b + z^2 c`
499        ring.add_assign(a, s1); // this is now `a + b + c`
500        hom.mul_assign_ref_map(b, t); // this is now `t (a + z^2 b + z c)`
501        hom.mul_assign_ref_map(c, t_sqr); // this is now `t^2 (a + z b + z^2 c`
502    }
503
504    ///
505    /// Possibly pre-processes elements before the FFT starts. Here you can bring ring element 
506    /// into a certain form, and assume during [`CooleyTukeyRadix3Butterfly::butterfly()`]
507    /// that the inputs are in this form.
508    /// 
509    #[inline(always)]
510    default fn prepare_for_fft(&self, _value: &mut Self::Element) {}
511    
512    ///
513    /// Possibly pre-processes elements before the inverse FFT starts. Here you can bring ring element
514    /// into a certain form, and assume during [`CooleyTukeyRadix3Butterfly::inv_butterfly()`]
515    /// that the inputs are in this form.
516    /// 
517    #[inline(always)]
518    default fn prepare_for_inv_fft(&self, _value: &mut Self::Element) {}
519}
520
521impl<H, A> FFTErrorEstimate for CooleyTukeyRadix3FFT<Complex64Base, Complex64Base, H, A> 
522    where H: Homomorphism<Complex64Base, Complex64Base>,
523        A: Allocator
524{
525    fn expected_absolute_error(&self, input_bound: f64, input_error: f64) -> f64 {
526        // the butterfly performs two multiplications with roots of unity, and then two additions
527        let multiply_absolute_error = 2. * input_bound * root_of_unity_error() + input_bound * f64::EPSILON;
528        let addition_absolute_error = 2. * input_bound * f64::EPSILON;
529        let butterfly_absolute_error = multiply_absolute_error + addition_absolute_error;
530        // the operator inf-norm of the FFT is its length
531        return 2. * self.len() as f64 * butterfly_absolute_error + self.len() as f64 * input_error;
532    }
533}
534
535#[cfg(test)]
536use std::array::from_fn;
537#[cfg(test)]
538use crate::rings::finite::FiniteRingStore;
539#[cfg(test)]
540use crate::rings::zn::zn_64::*;
541#[cfg(test)]
542use crate::rings::zn::zn_static::Fp;
543
544#[test]
545fn test_radix3_butterflies() {
546    let log3_n = 3;
547    let ring = Zn::new(109);
548    let ring_fastmul = ZnFastmul::new(ring).unwrap();
549    let int_hom = ring.int_hom();
550    let i = |x| int_hom.map(x);
551    let zeta = i(97);
552    let zeta_inv = ring.invert(&zeta).unwrap();
553    let fft = CooleyTukeyRadix3FFT::new_with_hom(ring.into_can_hom(ring_fastmul).ok().unwrap(), ring_fastmul.coerce(&ring, zeta), log3_n);
554
555    const LEN: usize = 27;
556    let data: [_; LEN] = from_fn(|j| i(j as i32));
557    let expected_std_order = |step: usize, group_idx: usize, value_idx: usize| ring.sum(
558        (0..ZZ.pow(3, step)).map(|k| ring.mul(
559            ring.pow(zeta_inv, value_idx * (k * ZZ.pow(3, log3_n - step)) as usize),
560            data[group_idx + (k * ZZ.pow(3, log3_n - step)) as usize]
561        ))
562    );
563    let expected_threeadic_reverse = |step: usize| from_fn(|i| expected_std_order(
564        step,
565        i % ZZ.pow(3, log3_n - step) as usize,
566        threeadic_reverse(i / ZZ.pow(3, log3_n - step) as usize, step)
567    ));
568    let begin = expected_threeadic_reverse(0);
569    for (a, e) in data.iter().zip(begin.iter()) {
570        assert_el_eq!(ring, a, e);
571    }
572
573    let mut actual = data;
574    for i in 0..log3_n {
575        fft.butterfly_step_main::<false>(&mut actual, i);
576        let expected: [ZnEl; LEN] = expected_threeadic_reverse(i + 1);
577        for (a, e) in actual.iter().zip(expected.iter()) {
578            assert_el_eq!(ring, a, e);
579        }
580    }
581}
582
583#[test]
584fn test_radix3_inv_fft() {
585    let log3_n = 3;
586    let ring = Zn::new(109);
587    let ring_fastmul = ZnFastmul::new(ring).unwrap();
588    let zeta = ring.int_hom().map(97);
589    let fft = CooleyTukeyRadix3FFT::new_with_hom(ring.into_can_hom(ring_fastmul).ok().unwrap(), ring_fastmul.coerce(&ring, zeta), log3_n);
590
591    let data = (0..ZZ.pow(3, log3_n)).map(|x| ring.int_hom().map(x as i32)).collect::<Vec<_>>();
592    let mut actual = data.clone();
593    fft.unordered_fft(&mut actual, &ring);
594    fft.unordered_inv_fft(&mut actual, &ring);
595
596    for i in 0..data.len() {
597        assert_el_eq!(&ring, &data[i], &actual[i]);
598    }
599}
600
601#[test]
602fn test_size_1_fft() {
603    let ring = Fp::<17>::RING;
604    let fft = CooleyTukeyRadix3FFT::for_zn(&ring, 0).unwrap().change_ring(ring.identity()).0;
605    let values: [u64; 1] = [3];
606    let mut work = values;
607    fft.unordered_fft(&mut work, ring);
608    assert_eq!(&work, &values);
609    fft.unordered_inv_fft(&mut work, ring);
610    assert_eq!(&work, &values);
611    assert_eq!(0, fft.unordered_fft_permutation(0));
612    assert_eq!(0, fft.unordered_fft_permutation_inv(0));
613}
614
615#[cfg(any(test, feature = "generic_tests"))]
616pub mod generic_tests {
617    use super::*;
618
619    pub fn test_cooley_tuckey_radix3_butterfly<R: RingStore, S: RingStore, I: Iterator<Item = El<R>>>(ring: R, base: S, edge_case_elements: I, test_zeta: &El<S>, test_twiddle: &El<S>)
620        where R::Type: CanHomFrom<S::Type>,
621            S::Type: DivisibilityRing
622    {
623        assert!(base.is_zero(&base.sum([base.one(), base.clone_el(&test_zeta), base.pow(base.clone_el(&test_zeta), 2)])));
624        let test_inv_twiddle = base.invert(&test_twiddle).unwrap();
625        let elements = edge_case_elements.collect::<Vec<_>>();
626        let hom = ring.can_hom(&base).unwrap();
627
628        for a in &elements {
629            for b in &elements {
630                for c in &elements {
631                    
632                    let [mut x, mut y, mut z] = [ring.clone_el(a), ring.clone_el(b), ring.clone_el(c)];
633                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::prepare_for_fft(ring.get_ring(), &mut x);
634                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::prepare_for_fft(ring.get_ring(), &mut y);
635                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::prepare_for_fft(ring.get_ring(), &mut z);
636                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::butterfly(
637                        &hom, 
638                        &mut x, 
639                        &mut y, 
640                        &mut z, 
641                        &test_zeta, 
642                        &test_twiddle, 
643                        &base.pow(base.mul_ref(&test_twiddle, &test_zeta), 2)
644                    );
645                    let mut t = hom.map_ref(&test_twiddle);
646                    assert_el_eq!(ring, ring.add_ref_fst(a, ring.mul_ref_snd(ring.add_ref_fst(b, ring.mul_ref(c, &t)), &t)), &x);
647                    ring.mul_assign(&mut t, hom.map_ref(&test_zeta));
648                    assert_el_eq!(ring, ring.add_ref_fst(a, ring.mul_ref_snd(ring.add_ref_fst(b, ring.mul_ref(c, &t)), &t)), &y);
649                    ring.mul_assign(&mut t, hom.map_ref(&test_zeta));
650                    assert_el_eq!(ring, ring.add_ref_fst(a, ring.mul_ref_snd(ring.add_ref_fst(b, ring.mul_ref(c, &t)), &t)), &z);
651                    
652                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::prepare_for_inv_fft(ring.get_ring(), &mut x);
653                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::prepare_for_inv_fft(ring.get_ring(), &mut y);
654                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::prepare_for_inv_fft(ring.get_ring(), &mut z);
655                    <R::Type as CooleyTukeyRadix3Butterfly<S::Type>>::inv_butterfly(
656                        &hom, 
657                        &mut x, 
658                        &mut y, 
659                        &mut z, 
660                        &test_zeta, 
661                        &test_inv_twiddle, 
662                        &base.pow(base.clone_el(&test_inv_twiddle), 2)
663                    );
664                    assert_el_eq!(ring, ring.int_hom().mul_ref_fst_map(a, 3), &x);
665                    assert_el_eq!(ring, ring.int_hom().mul_ref_fst_map(b, 3), &y);
666                    assert_el_eq!(ring, ring.int_hom().mul_ref_fst_map(c, 3), &z);
667                }
668            }
669        }
670    }
671}
672
673#[test]
674fn test_butterfly() {
675    let ring = Fp::<109>::RING;
676    generic_tests::test_cooley_tuckey_radix3_butterfly(
677        ring,
678        ring,
679        ring.elements().step_by(10),
680        &63,
681        &97,
682    );
683}
feanor_math/algorithms/fft/radix3.rs

feanor_math/algorithms/fft/
radix3.rs
