1
2use std::alloc::Allocator;
3
4use crate::seq::subvector::SubvectorView;
5use crate::ring::*;
6use crate::homomorphism::*;
7use crate::algorithms::fft::*;
8use crate::algorithms::unity_root::is_prim_root_of_unity;
9use crate::algorithms::fft::complex_fft::*;
10use crate::rings::float_complex::*;
11use crate::divisibility::DivisibilityRing;
12use crate::algorithms::fft::radix3::CooleyTukeyRadix3FFT;
13use crate::algorithms::fft::cooley_tuckey::CooleyTuckeyFFT;
14
15#[stability::unstable(feature = "enable")]
20pub struct GeneralCooleyTukeyFFT<R_main, R_twiddle, H, T1, T2>
21 where R_main: ?Sized + RingBase,
22 R_twiddle: ?Sized + RingBase,
23 H: Homomorphism<R_twiddle, R_main>,
24 T1: FFTAlgorithm<R_main>,
25 T2: FFTAlgorithm<R_main>
26{
27 twiddle_factors: Vec<R_twiddle::Element>,
28 inv_twiddle_factors: Vec<R_twiddle::Element>,
29 left_table: T1,
30 right_table: T2,
31 root_of_unity: R_main::Element,
32 root_of_unity_twiddle: R_twiddle::Element,
33 hom: H
34}
35
36impl<R, T1, T2> GeneralCooleyTukeyFFT<R::Type, R::Type, Identity<R>, T1, T2>
37 where R: RingStore,
38 T1: FFTAlgorithm<R::Type>,
39 T2: FFTAlgorithm<R::Type>
40{
41 #[stability::unstable(feature = "enable")]
50 pub fn new_with_pows<F>(ring: R, root_of_unity_pows: F, left_table: T1, right_table: T2) -> Self
51 where F: FnMut(i64) -> El<R>
52 {
53 Self::new_with_pows_with_hom(ring.into_identity(), root_of_unity_pows, left_table, right_table)
54 }
55
56 #[stability::unstable(feature = "enable")]
68 pub fn new(ring: R, root_of_unity: El<R>, left_table: T1, right_table: T2) -> Self {
69 Self::new_with_hom(ring.into_identity(), root_of_unity, left_table, right_table)
70 }
71}
72
73impl<R_main, R_twiddle, H, A1, A2> GeneralCooleyTukeyFFT<R_main, R_twiddle, H, CooleyTukeyRadix3FFT<R_main, R_twiddle, H, A1>, CooleyTuckeyFFT<R_main, R_twiddle, H, A2>>
74 where R_main: ?Sized + RingBase,
75 R_twiddle: ?Sized + RingBase + DivisibilityRing,
76 H: Homomorphism<R_twiddle, R_main>,
77 A1: Allocator + Clone,
78 A2: Allocator + Clone
79{
80 #[stability::unstable(feature = "enable")]
89 pub fn change_ring<R_new: ?Sized + RingBase, H_new: Clone + Homomorphism<R_twiddle, R_new>>(self, new_hom: H_new) -> (GeneralCooleyTukeyFFT<R_new, R_twiddle, H_new, CooleyTukeyRadix3FFT<R_new, R_twiddle, H_new, A1>, CooleyTuckeyFFT<R_new, R_twiddle, H_new, A2>>, H) {
90 let ring = new_hom.codomain();
91 let root_of_unity = new_hom.map_ref(&self.root_of_unity_twiddle);
92 assert!(ring.is_commutative());
93 assert!(ring.get_ring().is_approximate() || is_prim_root_of_unity(&ring, &root_of_unity, self.len()));
94
95 return (
96 GeneralCooleyTukeyFFT {
97 twiddle_factors: self.twiddle_factors,
98 left_table: self.left_table.change_ring(new_hom.clone()).0,
99 right_table: self.right_table.change_ring(new_hom.clone()).0,
100 inv_twiddle_factors: self.inv_twiddle_factors,
101 root_of_unity: root_of_unity,
102 root_of_unity_twiddle: self.root_of_unity_twiddle,
103 hom: new_hom,
104 },
105 self.hom
106 );
107 }
108}
109
110impl<R_main, R_twiddle, H, T1, T2> GeneralCooleyTukeyFFT<R_main, R_twiddle, H, T1, T2>
111 where R_main: ?Sized + RingBase,
112 R_twiddle: ?Sized + RingBase,
113 H: Homomorphism<R_twiddle, R_main>,
114 T1: FFTAlgorithm<R_main>,
115 T2: FFTAlgorithm<R_main>
116{
117 #[stability::unstable(feature = "enable")]
131 pub fn new_with_pows_with_hom<F>(hom: H, root_of_unity_pows: F, left_table: T1, right_table: T2) -> Self
132 where F: FnMut(i64) -> R_twiddle::Element
133 {
134 Self::create(hom, root_of_unity_pows, left_table, right_table)
135 }
136
137 #[stability::unstable(feature = "enable")]
141 pub fn create<F>(hom: H, mut root_of_unity_pows: F, left_table: T1, right_table: T2) -> Self
142 where F: FnMut(i64) -> R_twiddle::Element
143 {
144 let ring = hom.codomain();
145
146 assert!(ring.is_commutative());
147 assert!(ring.get_ring().is_approximate() || is_prim_root_of_unity(ring, &hom.map(root_of_unity_pows(1)), left_table.len() * right_table.len()));
148 assert!(ring.get_ring().is_approximate() || ring.eq_el(&hom.map(root_of_unity_pows(right_table.len().try_into().unwrap())), left_table.root_of_unity(ring)));
149 assert!(ring.get_ring().is_approximate() || ring.eq_el(&hom.map(root_of_unity_pows(left_table.len().try_into().unwrap())), right_table.root_of_unity(ring)));
150
151 let root_of_unity = root_of_unity_pows(1);
152 let inv_twiddle_factors = Self::create_twiddle_factors(|i| root_of_unity_pows(-i), &left_table, &right_table);
153 let twiddle_factors = Self::create_twiddle_factors(root_of_unity_pows, &left_table, &right_table);
154
155 GeneralCooleyTukeyFFT {
156 twiddle_factors: twiddle_factors,
157 inv_twiddle_factors: inv_twiddle_factors,
158 left_table: left_table,
159 right_table: right_table,
160 root_of_unity: hom.map_ref(&root_of_unity),
161 root_of_unity_twiddle: root_of_unity,
162 hom: hom
163 }
164 }
165
166 #[stability::unstable(feature = "enable")]
170 pub fn left_fft_table(&self) -> &T1 {
171 &self.left_table
172 }
173
174 #[stability::unstable(feature = "enable")]
178 pub fn right_fft_table(&self) -> &T2 {
179 &self.right_table
180 }
181
182 #[stability::unstable(feature = "enable")]
187 pub fn hom<'a>(&'a self) -> &'a H {
188 &self.hom
189 }
190
191 #[stability::unstable(feature = "enable")]
208 pub fn new_with_hom(hom: H, root_of_unity: R_twiddle::Element, left_table: T1, right_table: T2) -> Self {
209 let len = left_table.len() * right_table.len();
210 let root_of_unity_pows = |i: i64| if i >= 0 {
211 hom.domain().pow(hom.domain().clone_el(&root_of_unity), i.try_into().unwrap())
212 } else {
213 let len_i64: i64 = len.try_into().unwrap();
214 hom.domain().pow(hom.domain().clone_el(&root_of_unity), (len_i64 + (i % len_i64)).try_into().unwrap())
215 };
216 let result = GeneralCooleyTukeyFFT::create(&hom, root_of_unity_pows, left_table, right_table);
217 GeneralCooleyTukeyFFT {
218 twiddle_factors: result.twiddle_factors,
219 inv_twiddle_factors: result.inv_twiddle_factors,
220 left_table: result.left_table,
221 right_table: result.right_table,
222 root_of_unity: result.root_of_unity,
223 root_of_unity_twiddle: result.root_of_unity_twiddle,
224 hom: hom
225 }
226 }
227
228 fn create_twiddle_factors<F>(mut root_of_unity_pows: F, left_table: &T1, right_table: &T2) -> Vec<R_twiddle::Element>
229 where F: FnMut(i64) -> R_twiddle::Element
230 {
231 (0..(left_table.len() * right_table.len())).map(|i| {
232 let ri: i64 = (i % right_table.len()).try_into().unwrap();
233 let li = i / right_table.len();
234 return root_of_unity_pows(TryInto::<i64>::try_into(left_table.unordered_fft_permutation(li)).unwrap() * ri);
235 }).collect::<Vec<_>>()
236 }
237
238 #[stability::unstable(feature = "enable")]
242 pub fn ring<'a>(&'a self) -> &'a <H as Homomorphism<R_twiddle, R_main>>::CodomainStore {
243 self.hom.codomain()
244 }
245}
246
247impl<R_main, R_twiddle, H, T1, T2> PartialEq for GeneralCooleyTukeyFFT<R_main, R_twiddle, H, T1, T2>
248 where R_main: ?Sized + RingBase,
249 R_twiddle: ?Sized + RingBase,
250 H: Homomorphism<R_twiddle, R_main>,
251 T1: FFTAlgorithm<R_main> + PartialEq,
252 T2: FFTAlgorithm<R_main> + PartialEq
253{
254 fn eq(&self, other: &Self) -> bool {
255 self.ring().get_ring() == other.ring().get_ring() &&
256 self.left_table == other.left_table &&
257 self.right_table == other.right_table &&
258 self.ring().eq_el(self.root_of_unity(self.ring()), other.root_of_unity(self.ring()))
259 }
260}
261
262impl<R_main, R_twiddle, H, T1, T2> FFTAlgorithm<R_main> for GeneralCooleyTukeyFFT<R_main, R_twiddle, H, T1, T2>
263 where R_main: ?Sized + RingBase,
264 R_twiddle: ?Sized + RingBase,
265 H: Homomorphism<R_twiddle, R_main>,
266 T1: FFTAlgorithm<R_main>,
267 T2: FFTAlgorithm<R_main>
268{
269 fn len(&self) -> usize {
270 self.left_table.len() * self.right_table.len()
271 }
272
273 fn root_of_unity<S: RingStore<Type = R_main> + Copy>(&self, ring: S) -> &R_main::Element {
274 assert!(self.ring().get_ring() == ring.get_ring(), "unsupported ring");
275 &self.root_of_unity
276 }
277
278 fn unordered_fft<V, S>(&self, mut values: V, ring: S)
279 where V: SwappableVectorViewMut<<R_main as RingBase>::Element>,
280 S: RingStore<Type = R_main> + Copy
281 {
282 assert!(self.ring().get_ring() == ring.get_ring(), "unsupported ring");
283 if self.left_table.len() > 1 {
284 for i in 0..self.right_table.len() {
285 let mut v = SubvectorView::new(&mut values).restrict(i..).step_by_view(self.right_table.len());
286 self.left_table.unordered_fft(&mut v, ring);
287 }
288 for i in 0..self.len() {
289 self.hom.mul_assign_ref_map(values.at_mut(i), self.inv_twiddle_factors.at(i));
290 }
291 }
292 for i in 0..self.left_table.len() {
293 let mut v = SubvectorView::new(&mut values).restrict((i * self.right_table.len())..((i + 1) * self.right_table.len()));
294 self.right_table.unordered_fft(&mut v, ring);
295 }
296 }
297
298 fn unordered_inv_fft<V, S>(&self, mut values: V, ring: S)
299 where V: SwappableVectorViewMut<<R_main as RingBase>::Element>,
300 S: RingStore<Type = R_main> + Copy
301 {
302 assert!(self.ring().get_ring() == ring.get_ring(), "unsupported ring");
303 for i in 0..self.left_table.len() {
304 let mut v = SubvectorView::new(&mut values).restrict((i * self.right_table.len())..((i + 1) * self.right_table.len()));
305 self.right_table.unordered_inv_fft(&mut v, ring);
306 }
307 if self.left_table.len() > 1 {
308 for i in 0..self.len() {
309 self.hom.mul_assign_ref_map(values.at_mut(i), self.twiddle_factors.at(i));
310 debug_assert!(self.ring().get_ring().is_approximate() || self.hom.domain().is_one(&self.hom.domain().mul_ref(self.twiddle_factors.at(i), self.inv_twiddle_factors.at(i))));
311 }
312 for i in 0..self.right_table.len() {
313 let mut v = SubvectorView::new(&mut values).restrict(i..).step_by_view(self.right_table.len());
314 self.left_table.unordered_inv_fft(&mut v, ring);
315 }
316 }
317 }
318
319 fn unordered_fft_permutation(&self, i: usize) -> usize {
320 assert!(i < self.len());
321 self.left_table.unordered_fft_permutation(i / self.right_table.len()) + self.left_table.len() * self.right_table.unordered_fft_permutation(i % self.right_table.len())
322 }
323
324 fn unordered_fft_permutation_inv(&self, i: usize) -> usize {
325 assert!(i < self.len());
326 self.left_table.unordered_fft_permutation_inv(i % self.left_table.len()) * self.right_table.len() + self.right_table.unordered_fft_permutation_inv(i / self.left_table.len())
327 }
328}
329
330impl<H, T1, T2> FFTErrorEstimate for GeneralCooleyTukeyFFT<Complex64Base, Complex64Base, H, T1, T2>
331 where H: Homomorphism<Complex64Base, Complex64Base>,
332 T1: FFTAlgorithm<Complex64Base> + FFTErrorEstimate,
333 T2: FFTAlgorithm<Complex64Base> + FFTErrorEstimate
334{
335 fn expected_absolute_error(&self, input_bound: f64, input_error: f64) -> f64 {
336 let error_after_first_fft = self.left_table.expected_absolute_error(input_bound, input_error);
337 let new_input_bound = self.left_table.len() as f64 * input_bound;
338 let error_after_twiddling = error_after_first_fft + new_input_bound * (root_of_unity_error() + f64::EPSILON);
339 return self.right_table.expected_absolute_error(new_input_bound, error_after_twiddling);
340 }
341}
342
343#[cfg(test)]
344use crate::rings::zn::zn_static::{Zn, Fp};
345#[cfg(test)]
346use crate::algorithms::unity_root::*;
347#[cfg(test)]
348use crate::rings::zn::zn_64;
349#[cfg(test)]
350use crate::rings::zn::ZnRingStore;
351#[cfg(test)]
352use std::alloc::Global;
353#[cfg(test)]
354use crate::algorithms::fft::bluestein::BluesteinFFT;
355
356#[test]
357fn test_fft_basic() {
358 let ring = Zn::<97>::RING;
359 let z = ring.int_hom().map(39);
360 let fft = GeneralCooleyTukeyFFT::new(ring, ring.pow(z, 16),
361 CooleyTuckeyFFT::new(ring, ring.pow(z, 48), 1),
362 BluesteinFFT::new(ring, ring.pow(z, 16), ring.pow(z, 12), 3, 3, Global),
363 );
364 let mut values = [1, 0, 0, 1, 0, 1];
365 let expected = [3, 62, 63, 96, 37, 36];
366 let mut permuted_expected = [0; 6];
367 for i in 0..6 {
368 permuted_expected[i] = expected[fft.unordered_fft_permutation(i)];
369 }
370
371 fft.unordered_fft(&mut values, ring);
372 assert_eq!(values, permuted_expected);
373}
374
375#[test]
376fn test_fft_long() {
377 let ring = Fp::<97>::RING;
378 let z = ring.int_hom().map(39);
379 let fft = GeneralCooleyTukeyFFT::new(ring, ring.pow(z, 4),
380 CooleyTuckeyFFT::new(ring, ring.pow(z, 12), 3),
381 BluesteinFFT::new(ring, ring.pow(z, 16), ring.pow(z, 12), 3, 3, Global),
382 );
383 let mut values = [1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 2, 2, 0, 2, 0, 1, 2, 3, 4];
384 let expected = [26, 0, 75, 47, 41, 31, 28, 62, 39, 93, 53, 27, 0, 54, 74, 61, 65, 81, 63, 38, 53, 94, 89, 91];
385 let mut permuted_expected = [0; 24];
386 for i in 0..24 {
387 permuted_expected[i] = expected[fft.unordered_fft_permutation(i)];
388 }
389
390 fft.unordered_fft(&mut values, ring);
391 assert_eq!(values, permuted_expected);
392}
393
394#[test]
395fn test_fft_unordered() {
396 let ring = Fp::<1409>::RING;
397 let z = get_prim_root_of_unity(ring, 64 * 11).unwrap();
398 let fft = GeneralCooleyTukeyFFT::new(
399 ring,
400 ring.pow(z, 4),
401 CooleyTuckeyFFT::new(ring, ring.pow(z, 44), 4),
402 BluesteinFFT::new(ring, ring.pow(z, 32), ring.pow(z, 22), 11, 5, Global),
403 );
404 const LEN: usize = 16 * 11;
405 let mut values = [0; LEN];
406 for i in 0..LEN {
407 values[i] = ring.int_hom().map(i as i32);
408 }
409 let original = values;
410
411 fft.unordered_fft(&mut values, ring);
412
413 let mut ordered_fft = [0; LEN];
414 for i in 0..LEN {
415 ordered_fft[fft.unordered_fft_permutation(i)] = values[i];
416 }
417
418 fft.unordered_inv_fft(&mut values, ring);
419 assert_eq!(values, original);
420
421 fft.inv_fft(&mut ordered_fft, ring);
422 assert_eq!(ordered_fft, original);
423}
424
425
426#[test]
427fn test_unordered_fft_permutation_inv() {
428 let ring = Fp::<1409>::RING;
429 let z = get_prim_root_of_unity(ring, 64 * 11).unwrap();
430 let fft = GeneralCooleyTukeyFFT::new(
431 ring,
432 ring.pow(z, 4),
433 CooleyTuckeyFFT::new(ring, ring.pow(z, 44), 4),
434 BluesteinFFT::new(ring, ring.pow(z, 32), ring.pow(z, 22), 11, 5, Global),
435 );
436 for i in 0..(16 * 11) {
437 assert_eq!(fft.unordered_fft_permutation_inv(fft.unordered_fft_permutation(i)), i);
438 assert_eq!(fft.unordered_fft_permutation(fft.unordered_fft_permutation_inv(i)), i);
439 }
440
441 let fft = GeneralCooleyTukeyFFT::new(
442 ring,
443 ring.pow(z, 4),
444 BluesteinFFT::new(ring, ring.pow(z, 32), ring.pow(z, 22), 11, 5, Global),
445 CooleyTuckeyFFT::new(ring, ring.pow(z, 44), 4),
446 );
447 for i in 0..(16 * 11) {
448 assert_eq!(fft.unordered_fft_permutation_inv(fft.unordered_fft_permutation(i)), i);
449 assert_eq!(fft.unordered_fft_permutation(fft.unordered_fft_permutation_inv(i)), i);
450 }
451}
452
453#[test]
454fn test_inv_fft() {
455 let ring = Fp::<97>::RING;
456 let z = ring.int_hom().map(39);
457 let fft = GeneralCooleyTukeyFFT::new(ring, ring.pow(z, 16),
458 CooleyTuckeyFFT::new(ring, ring.pow(z, 16 * 3), 1),
459 BluesteinFFT::new(ring, ring.pow(z, 16), ring.pow(z, 12), 3, 3, Global),
460 );
461 let mut values = [3, 62, 63, 96, 37, 36];
462 let expected = [1, 0, 0, 1, 0, 1];
463
464 fft.inv_fft(&mut values, ring);
465 assert_eq!(values, expected);
466}
467
468#[test]
469fn test_approximate_fft() {
470 let CC = Complex64::RING;
471 for (p, log2_n) in [(5, 3), (53, 5), (101, 8), (503, 10)] {
472 let fft = GeneralCooleyTukeyFFT::new_with_pows(
473 CC,
474 |i| CC.root_of_unity(i, TryInto::<i64>::try_into(p).unwrap() << log2_n),
475 BluesteinFFT::for_complex(CC, p, Global),
476 CooleyTuckeyFFT::for_complex(CC, log2_n)
477 );
478 let mut array = (0..(p << log2_n)).map(|i| CC.root_of_unity(i.try_into().unwrap(), TryInto::<i64>::try_into(p).unwrap() << log2_n)).collect::<Vec<_>>();
479 fft.fft(&mut array, CC);
480 let err = fft.expected_absolute_error(1., 0.);
481 assert!(CC.is_absolute_approx_eq(array[0], CC.zero(), err));
482 assert!(CC.is_absolute_approx_eq(array[1], CC.from_f64(fft.len() as f64), err));
483 for i in 2..fft.len() {
484 assert!(CC.is_absolute_approx_eq(array[i], CC.zero(), err));
485 }
486 }
487}
488
489#[cfg(test)]
490const BENCH_N1: usize = 31;
491#[cfg(test)]
492const BENCH_N2: usize = 601;
493
494#[bench]
495fn bench_factor_fft(bencher: &mut test::Bencher) {
496 let ring = zn_64::Zn::new(1602564097);
497 let fastmul_ring = zn_64::ZnFastmul::new(ring).unwrap();
498 let embedding = ring.can_hom(&fastmul_ring).unwrap();
499 let ring_as_field = ring.as_field().ok().unwrap();
500 let root_of_unity = fastmul_ring.coerce(&ring, ring_as_field.get_ring().unwrap_element(get_prim_root_of_unity(&ring_as_field, 2 * 31 * 601).unwrap()));
501 let fastmul_ring_as_field = fastmul_ring.as_field().ok().unwrap();
502 let fft = GeneralCooleyTukeyFFT::new_with_hom(
503 embedding.clone(),
504 fastmul_ring.pow(root_of_unity, 2),
505 BluesteinFFT::new_with_hom(embedding.clone(), fastmul_ring.pow(root_of_unity, BENCH_N1), fastmul_ring_as_field.get_ring().unwrap_element(get_prim_root_of_unity_pow2(&fastmul_ring_as_field, 11).unwrap()), BENCH_N2, 11, Global),
506 BluesteinFFT::new_with_hom(embedding, fastmul_ring.pow(root_of_unity, BENCH_N2), fastmul_ring_as_field.get_ring().unwrap_element(get_prim_root_of_unity_pow2(&fastmul_ring_as_field, 6).unwrap()), BENCH_N1, 6, Global),
507 );
508 let data = (0..(BENCH_N1 * BENCH_N2)).map(|i| ring.int_hom().map(i as i32)).collect::<Vec<_>>();
509 let mut copy = Vec::with_capacity(BENCH_N1 * BENCH_N2);
510 bencher.iter(|| {
511 copy.clear();
512 copy.extend(data.iter().map(|x| ring.clone_el(x)));
513 fft.unordered_fft(&mut copy[..], &ring);
514 fft.unordered_inv_fft(&mut copy[..], &ring);
515 assert_el_eq!(ring, copy[0], data[0]);
516 });
517}