singe_kernel/cpu/
fft.rs

1//! Small FFT helpers, interleaved complex transforms, and normalization.
2
3use num_complex::Complex32;
4
5#[derive(Clone, Copy, Debug, Eq, PartialEq)]
6pub enum FftDirection {
7    Forward,
8    Inverse,
9}
10
11pub fn c2c(input: &[Complex32], batch: usize, n: usize, direction: FftDirection) -> Vec<Complex32> {
12    let sign = match direction {
13        FftDirection::Forward => -1.0,
14        FftDirection::Inverse => 1.0,
15    };
16    let mut output = vec![Complex32::new(0.0, 0.0); batch * n];
17    for batch_index in 0..batch {
18        let base = batch_index * n;
19        for frequency in 0..n {
20            let mut real = 0.0;
21            let mut imag = 0.0;
22            for sample in 0..n {
23                let value = input[base + sample];
24                let angle =
25                    sign * 2.0 * std::f32::consts::PI * frequency as f32 * sample as f32 / n as f32;
26                real += value.re * angle.cos() - value.im * angle.sin();
27                imag += value.re * angle.sin() + value.im * angle.cos();
28            }
29            output[base + frequency] = Complex32::new(real, imag);
30        }
31    }
32    output
33}
34
35pub fn r2c(input: &[f32], batch: usize, n: usize) -> Vec<Complex32> {
36    let frequency_bins = n / 2 + 1;
37    let mut output = vec![Complex32::new(0.0, 0.0); batch * frequency_bins];
38    for batch_index in 0..batch {
39        let input_base = batch_index * n;
40        let output_base = batch_index * frequency_bins;
41        for frequency in 0..frequency_bins {
42            let mut real = 0.0;
43            let mut imag = 0.0;
44            for sample in 0..n {
45                let value = input[input_base + sample];
46                let angle =
47                    -2.0 * std::f32::consts::PI * frequency as f32 * sample as f32 / n as f32;
48                real += value * angle.cos();
49                imag += value * angle.sin();
50            }
51            output[output_base + frequency] = Complex32::new(real, imag);
52        }
53    }
54    output
55}
56
57pub fn c2r_unscaled(input: &[Complex32], batch: usize, n: usize) -> Vec<f32> {
58    let frequency_bins = n / 2 + 1;
59    let mut output = vec![0.0; batch * n];
60    for batch_index in 0..batch {
61        let input_base = batch_index * frequency_bins;
62        let output_base = batch_index * n;
63        for sample in 0..n {
64            let mut value = input[input_base].re;
65            let nyquist = input[input_base + n / 2].re;
66            value += nyquist * if sample % 2 == 0 { 1.0 } else { -1.0 };
67            for frequency in 1..n / 2 {
68                let bin = input[input_base + frequency];
69                let angle =
70                    2.0 * std::f32::consts::PI * frequency as f32 * sample as f32 / n as f32;
71                value += 2.0 * (bin.re * angle.cos() - bin.im * angle.sin());
72            }
73            output[output_base + sample] = value;
74        }
75    }
76    output
77}
78
79pub fn dft_c2c_f32_interleaved(input: &[f32], batch: usize, n: usize) -> Vec<f32> {
80    let mut output = vec![0.0; input.len()];
81    for batch_index in 0..batch {
82        for k in 0..n {
83            let mut real = 0.0f32;
84            let mut imag = 0.0f32;
85            for j in 0..n {
86                let angle = -2.0 * std::f32::consts::PI * (k * j) as f32 / n as f32;
87                let twiddle_real = angle.cos();
88                let twiddle_imag = angle.sin();
89                let input_offset = (batch_index * n + j) * 2;
90                let input_real = input[input_offset];
91                let input_imag = input[input_offset + 1];
92                real += input_real * twiddle_real - input_imag * twiddle_imag;
93                imag += input_real * twiddle_imag + input_imag * twiddle_real;
94            }
95            let output_offset = (batch_index * n + k) * 2;
96            output[output_offset] = real;
97            output[output_offset + 1] = imag;
98        }
99    }
100    output
101}
102
103pub fn twiddle_tables(n: usize) -> (Vec<f32>, Vec<f32>) {
104    let mut real = Vec::with_capacity(n);
105    let mut imag = Vec::with_capacity(n);
106    for phase in 0..n {
107        let angle = -2.0 * std::f32::consts::PI * phase as f32 / n as f32;
108        real.push(angle.cos());
109        imag.push(angle.sin());
110    }
111    (real, imag)
112}
singe_kernel/cpu/fft.rs

singe_kernel/cpu/
fft.rs
