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