numra_fft/
utils.rs

1//! FFT utility functions: frequencies, shifting, windowing.
2//!
3//! Author: Moussa Leblouba
4//! Date: 5 March 2026
5//! Modified: 2 May 2026
6
7use numra_core::Scalar;
8
9/// Window function types for spectral analysis.
10#[derive(Clone, Debug)]
11pub enum Window {
12    /// Rectangular (no windowing).
13    Rectangular,
14    /// Hann window: 0.5 * (1 - cos(2*pi*n/(N-1))).
15    Hann,
16    /// Hamming window: 0.54 - 0.46 * cos(2*pi*n/(N-1)).
17    Hamming,
18    /// Blackman window.
19    Blackman,
20    /// Kaiser window with shape parameter beta.
21    Kaiser(f64),
22    /// Flat-top window for accurate amplitude measurement.
23    FlatTop,
24}
25
26/// Generate window function coefficients.
27///
28/// Returns a vector of length `n` with the window weights.
29pub fn window_func<S: Scalar>(window: &Window, n: usize) -> Vec<S> {
30    if n == 0 {
31        return vec![];
32    }
33    if n == 1 {
34        return vec![S::ONE];
35    }
36
37    let nm1 = S::from_usize(n - 1);
38    let two_pi = S::TWO * S::PI;
39    let four_pi = S::from_f64(4.0) * S::PI;
40
41    match window {
42        Window::Rectangular => vec![S::ONE; n],
43        Window::Hann => (0..n)
44            .map(|i| S::HALF * (S::ONE - (two_pi * S::from_usize(i) / nm1).cos()))
45            .collect(),
46        Window::Hamming => {
47            let a0 = S::from_f64(0.54);
48            let a1 = S::from_f64(0.46);
49            (0..n)
50                .map(|i| a0 - a1 * (two_pi * S::from_usize(i) / nm1).cos())
51                .collect()
52        }
53        Window::Blackman => {
54            let a0 = S::from_f64(0.42);
55            let a1 = S::HALF;
56            let a2 = S::from_f64(0.08);
57            (0..n)
58                .map(|i| {
59                    let x = S::from_usize(i) / nm1;
60                    a0 - a1 * (two_pi * x).cos() + a2 * (four_pi * x).cos()
61                })
62                .collect()
63        }
64        Window::Kaiser(beta) => kaiser_window(S::from_f64(*beta), n),
65        Window::FlatTop => {
66            // Flat-top window coefficients (five-term cosine sum)
67            let a0 = S::from_f64(0.21557895);
68            let a1 = S::from_f64(0.41663158);
69            let a2 = S::from_f64(0.277263158);
70            let a3 = S::from_f64(0.083578947);
71            let a4 = S::from_f64(0.006947368);
72            let six_pi = S::from_f64(6.0) * S::PI;
73            let eight_pi = S::from_f64(8.0) * S::PI;
74            (0..n)
75                .map(|i| {
76                    let x = S::from_usize(i) / nm1;
77                    a0 - a1 * (two_pi * x).cos() + a2 * (four_pi * x).cos()
78                        - a3 * (six_pi * x).cos()
79                        + a4 * (eight_pi * x).cos()
80                })
81                .collect()
82        }
83    }
84}
85
86/// Kaiser window using the I0 Bessel function approximation.
87fn kaiser_window<S: Scalar>(beta: S, n: usize) -> Vec<S> {
88    let nm1 = S::from_usize(n - 1);
89    let denom = bessel_i0(beta);
90    (0..n)
91        .map(|i| {
92            let x = S::TWO * S::from_usize(i) / nm1 - S::ONE;
93            let arg = beta * (S::ONE - x * x).max(S::ZERO).sqrt();
94            bessel_i0(arg) / denom
95        })
96        .collect()
97}
98
99/// Modified Bessel function I0(x) via power series.
100fn bessel_i0<S: Scalar>(x: S) -> S {
101    let mut sum = S::ONE;
102    let mut term = S::ONE;
103    let x2_4 = x * x / S::from_f64(4.0);
104    for k in 1..50 {
105        let kf = S::from_usize(k);
106        term = term * x2_4 / (kf * kf);
107        sum += term;
108        if term.to_f64() < 1e-16 * sum.to_f64() {
109            break;
110        }
111    }
112    sum
113}
114
115/// Compute the DFT sample frequencies.
116///
117/// Returns the frequency bin centers in cycles per unit of sample spacing `dt`.
118/// For N samples, returns N frequency values.
119///
120/// f = [0, 1, ..., N/2-1, -N/2, ..., -1] / (dt * N)
121pub fn fftfreq<S: Scalar>(n: usize, dt: S) -> Vec<S> {
122    let val = S::ONE / (S::from_usize(n) * dt);
123    let half = n.div_ceil(2);
124
125    (0..n)
126        .map(|i| {
127            if i < half {
128                S::from_usize(i) * val
129            } else {
130                (S::from_usize(i) - S::from_usize(n)) * val
131            }
132        })
133        .collect()
134}
135
136/// Shift zero-frequency component to the center.
137///
138/// Rearranges output of `fft` so that the zero-frequency component is centered.
139pub fn fftshift<S: Scalar>(x: &mut [S]) {
140    let n = x.len();
141    if n <= 1 {
142        return;
143    }
144    let mid = n / 2;
145    x.rotate_left(mid);
146}
147
148/// Inverse of `fftshift`: move center back to beginning.
149pub fn ifftshift<S: Scalar>(x: &mut [S]) {
150    let n = x.len();
151    if n <= 1 {
152        return;
153    }
154    let mid = n.div_ceil(2);
155    x.rotate_left(mid);
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn test_fftfreq_even() {
164        let f: Vec<f64> = fftfreq(8, 1.0);
165        assert_eq!(f.len(), 8);
166        assert!((f[0] - 0.0).abs() < 1e-14);
167        assert!((f[1] - 0.125).abs() < 1e-14);
168        assert!((f[4] - (-0.5)).abs() < 1e-14);
169        assert!((f[7] - (-0.125)).abs() < 1e-14);
170    }
171
172    #[test]
173    fn test_fftfreq_odd() {
174        let f: Vec<f64> = fftfreq(5, 0.1);
175        // Frequencies: [0, 2, 4, -4, -2]
176        assert!((f[0] - 0.0).abs() < 1e-14);
177        assert!((f[1] - 2.0).abs() < 1e-14);
178        assert!((f[2] - 4.0).abs() < 1e-14);
179        assert!((f[3] - (-4.0)).abs() < 1e-14);
180        assert!((f[4] - (-2.0)).abs() < 1e-14);
181    }
182
183    #[test]
184    fn test_fftshift() {
185        let mut x: Vec<f64> = vec![0.0, 1.0, 2.0, 3.0, -4.0, -3.0, -2.0, -1.0];
186        fftshift(&mut x);
187        assert_eq!(x, vec![-4.0, -3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0]);
188    }
189
190    #[test]
191    fn test_fftshift_ifftshift_roundtrip() {
192        let original: Vec<f64> = vec![0.0, 1.0, 2.0, -3.0, -2.0, -1.0];
193        let mut x = original.clone();
194        fftshift(&mut x);
195        ifftshift(&mut x);
196        for (a, b) in original.iter().zip(x.iter()) {
197            assert!((a - b).abs() < 1e-14);
198        }
199    }
200
201    #[test]
202    fn test_window_rectangular() {
203        let w: Vec<f64> = window_func(&Window::Rectangular, 5);
204        assert!(w.iter().all(|&v| (v - 1.0).abs() < 1e-14));
205    }
206
207    #[test]
208    fn test_window_hann_endpoints() {
209        let w: Vec<f64> = window_func(&Window::Hann, 8);
210        assert!(w[0].abs() < 1e-14); // Starts at 0
211        assert!(w[7].abs() < 1e-14); // Ends at 0
212        assert!((w[4] - 1.0).abs() < 0.1); // Peak near center
213    }
214
215    #[test]
216    fn test_window_hamming() {
217        let w: Vec<f64> = window_func(&Window::Hamming, 8);
218        // Hamming doesn't go to zero at edges
219        assert!(w[0] > 0.05);
220        assert!(w[7] > 0.05);
221    }
222
223    #[test]
224    fn test_window_blackman() {
225        let w: Vec<f64> = window_func(&Window::Blackman, 16);
226        assert!(w[0].abs() < 1e-10); // Very small at edges
227        assert!(w[8] > 0.9); // Peak near center
228    }
229
230    #[test]
231    fn test_window_kaiser() {
232        let w: Vec<f64> = window_func(&Window::Kaiser(5.0), 16);
233        assert!(w.iter().all(|&v| v >= 0.0 && v <= 1.01));
234        // Kaiser is symmetric
235        for i in 0..8 {
236            assert!((w[i] - w[15 - i]).abs() < 1e-12);
237        }
238    }
239
240    #[test]
241    fn test_window_empty() {
242        assert!(window_func::<f64>(&Window::Hann, 0).is_empty());
243        assert_eq!(window_func::<f64>(&Window::Hann, 1), vec![1.0]);
244    }
245}
numra_fft/utils.rs

numra_fft/
utils.rs
