1use std::fmt::Debug;
  2
  3use rand_distr::num_traits::{Float, FloatConst};
  4
  5use crate::{
  6    alloc_aligned,
  7    reference::fft64::reim::{ReimDFTExecute, fft_ref, frac_rev_bits},
  8};
  9
 10pub struct ReimFFTRef;
 11
 12impl ReimDFTExecute<ReimFFTTable<f64>, f64> for ReimFFTRef {
 13    fn reim_dft_execute(table: &ReimFFTTable<f64>, data: &mut [f64]) {
 14        fft_ref(table.m, &table.omg, data);
 15    }
 16}
 17
 18pub struct ReimFFTTable<R: Float + FloatConst + Debug> {
 19    m: usize,
 20    omg: Vec<R>,
 21}
 22
 23impl<R: Float + FloatConst + Debug + 'static> ReimFFTTable<R> {
 24    pub fn new(m: usize) -> Self {
 25        assert!(m & (m - 1) == 0, "m must be a power of two but is {m}");
 26        let mut omg: Vec<R> = alloc_aligned::<R>(2 * m);
 27
 28        let quarter: R = R::from(1. / 4.).unwrap();
 29
 30        if m <= 16 {
 31            match m {
 32                1 => {}
 33                2 => {
 34                    fill_fft2_omegas(quarter, &mut omg, 0);
 35                }
 36                4 => {
 37                    fill_fft4_omegas(quarter, &mut omg, 0);
 38                }
 39                8 => {
 40                    fill_fft8_omegas(quarter, &mut omg, 0);
 41                }
 42                16 => {
 43                    fill_fft16_omegas(quarter, &mut omg, 0);
 44                }
 45                _ => {}
 46            }
 47        } else if m <= 2048 {
 48            fill_fft_bfs_16_omegas(m, quarter, &mut omg, 0);
 49        } else {
 50            fill_fft_rec_16_omegas(m, quarter, &mut omg, 0);
 51        }
 52
 53        Self { m, omg }
 54    }
 55
 56    pub fn m(&self) -> usize {
 57        self.m
 58    }
 59
 60    pub fn omg(&self) -> &[R] {
 61        &self.omg
 62    }
 63}
 64
 65#[inline(always)]
 66fn fill_fft2_omegas<R: Float + FloatConst>(j: R, omg: &mut [R], pos: usize) -> usize {
 67    let omg_pos: &mut [R] = &mut omg[pos..];
 68    assert!(omg_pos.len() >= 2);
 69    let angle: R = j / R::from(2).unwrap();
 70    let two_pi: R = R::from(2).unwrap() * R::PI();
 71    omg_pos[0] = R::cos(two_pi * angle);
 72    omg_pos[1] = R::sin(two_pi * angle);
 73    pos + 2
 74}
 75
 76#[inline(always)]
 77fn fill_fft4_omegas<R: Float + FloatConst>(j: R, omg: &mut [R], pos: usize) -> usize {
 78    let omg_pos: &mut [R] = &mut omg[pos..];
 79    assert!(omg_pos.len() >= 4);
 80    let angle_1: R = j / R::from(2).unwrap();
 81    let angle_2: R = j / R::from(4).unwrap();
 82    let two_pi: R = R::from(2).unwrap() * R::PI();
 83    omg_pos[0] = R::cos(two_pi * angle_1);
 84    omg_pos[1] = R::sin(two_pi * angle_1);
 85    omg_pos[2] = R::cos(two_pi * angle_2);
 86    omg_pos[3] = R::sin(two_pi * angle_2);
 87    pos + 4
 88}
 89
 90#[inline(always)]
 91fn fill_fft8_omegas<R: Float + FloatConst>(j: R, omg: &mut [R], pos: usize) -> usize {
 92    let omg_pos: &mut [R] = &mut omg[pos..];
 93    assert!(omg_pos.len() >= 8);
 94    let _8th: R = R::from(1. / 8.).unwrap();
 95    let angle_1: R = j / R::from(2).unwrap();
 96    let angle_2: R = j / R::from(4).unwrap();
 97    let angle_4: R = j / R::from(8).unwrap();
 98    let two_pi: R = R::from(2).unwrap() * R::PI();
 99    omg_pos[0] = R::cos(two_pi * angle_1);
100    omg_pos[1] = R::sin(two_pi * angle_1);
101    omg_pos[2] = R::cos(two_pi * angle_2);
102    omg_pos[3] = R::sin(two_pi * angle_2);
103    omg_pos[4] = R::cos(two_pi * angle_4);
104    omg_pos[5] = R::cos(two_pi * (angle_4 + _8th));
105    omg_pos[6] = R::sin(two_pi * angle_4);
106    omg_pos[7] = R::sin(two_pi * (angle_4 + _8th));
107    pos + 8
108}
109
110#[inline(always)]
111fn fill_fft16_omegas<R: Float + FloatConst>(j: R, omg: &mut [R], pos: usize) -> usize {
112    let omg_pos: &mut [R] = &mut omg[pos..];
113    assert!(omg_pos.len() >= 16);
114    let _8th: R = R::from(1. / 8.).unwrap();
115    let _16th: R = R::from(1. / 16.).unwrap();
116    let angle_1: R = j / R::from(2).unwrap();
117    let angle_2: R = j / R::from(4).unwrap();
118    let angle_4: R = j / R::from(8).unwrap();
119    let angle_8: R = j / R::from(16).unwrap();
120    let two_pi: R = R::from(2).unwrap() * R::PI();
121    omg_pos[0] = R::cos(two_pi * angle_1);
122    omg_pos[1] = R::sin(two_pi * angle_1);
123    omg_pos[2] = R::cos(two_pi * angle_2);
124    omg_pos[3] = R::sin(two_pi * angle_2);
125    omg_pos[4] = R::cos(two_pi * angle_4);
126    omg_pos[5] = R::sin(two_pi * angle_4);
127    omg_pos[6] = R::cos(two_pi * (angle_4 + _8th));
128    omg_pos[7] = R::sin(two_pi * (angle_4 + _8th));
129    omg_pos[8] = R::cos(two_pi * angle_8);
130    omg_pos[9] = R::cos(two_pi * (angle_8 + _8th));
131    omg_pos[10] = R::cos(two_pi * (angle_8 + _16th));
132    omg_pos[11] = R::cos(two_pi * (angle_8 + _8th + _16th));
133    omg_pos[12] = R::sin(two_pi * angle_8);
134    omg_pos[13] = R::sin(two_pi * (angle_8 + _8th));
135    omg_pos[14] = R::sin(two_pi * (angle_8 + _16th));
136    omg_pos[15] = R::sin(two_pi * (angle_8 + _8th + _16th));
137    pos + 16
138}
139
140#[inline(always)]
141fn fill_fft_bfs_16_omegas<R: Float + FloatConst>(m: usize, j: R, omg: &mut [R], mut pos: usize) -> usize {
142    let log_m: usize = (usize::BITS - (m - 1).leading_zeros()) as usize;
143    let mut mm: usize = m;
144    let mut jj: R = j;
145
146    let two_pi: R = R::from(2).unwrap() * R::PI();
147
148    if !log_m.is_multiple_of(2) {
149        let h = mm >> 1;
150        let j: R = jj * R::from(0.5).unwrap();
151        omg[pos] = R::cos(two_pi * j);
152        omg[pos + 1] = R::sin(two_pi * j);
153        pos += 2;
154        mm = h;
155        jj = j
156    }
157
158    while mm > 16 {
159        let h: usize = mm >> 2;
160        let j: R = jj * R::from(1. / 4.).unwrap();
161        for i in (0..m).step_by(mm) {
162            let rs_0 = j + frac_rev_bits::<R>(i / mm) * R::from(1. / 4.).unwrap();
163            let rs_1 = R::from(2).unwrap() * rs_0;
164            omg[pos] = R::cos(two_pi * rs_1);
165            omg[pos + 1] = R::sin(two_pi * rs_1);
166            omg[pos + 2] = R::cos(two_pi * rs_0);
167            omg[pos + 3] = R::sin(two_pi * rs_0);
168            pos += 4;
169        }
170        mm = h;
171        jj = j;
172    }
173
174    for i in (0..m).step_by(16) {
175        let j = jj + frac_rev_bits(i >> 4);
176        fill_fft16_omegas(j, omg, pos);
177        pos += 16
178    }
179
180    pos
181}
182
183#[inline(always)]
184fn fill_fft_rec_16_omegas<R: Float + FloatConst>(m: usize, j: R, omg: &mut [R], mut pos: usize) -> usize {
185    if m <= 2048 {
186        return fill_fft_bfs_16_omegas(m, j, omg, pos);
187    }
188    let h: usize = m >> 1;
189    let s: R = j * R::from(0.5).unwrap();
190    let _2pi = R::from(2).unwrap() * R::PI();
191    omg[pos] = R::cos(_2pi * s);
192    omg[pos + 1] = R::sin(_2pi * s);
193    pos += 2;
194    pos = fill_fft_rec_16_omegas(h, s, omg, pos);
195    pos = fill_fft_rec_16_omegas(h, s + R::from(0.5).unwrap(), omg, pos);
196    pos
197}
198
199#[inline(always)]
200fn ctwiddle_ref(ra: &mut f64, ia: &mut f64, rb: &mut f64, ib: &mut f64, omg_re: f64, omg_im: f64) {
201    let dr: f64 = *rb * omg_re - *ib * omg_im;
202    let di: f64 = *rb * omg_im + *ib * omg_re;
203    *rb = *ra - dr;
204    *ib = *ia - di;
205    *ra += dr;
206    *ia += di;
207}