u_numflow/
fourier.rs

1//! Discrete Fourier transform.
2//!
3//! Provides the forward and inverse DFT of a complex sequence of **any
4//! length**, and a real-input convenience. A power-of-two length is
5//! transformed by the iterative radix-2 Cooley-Tukey algorithm; any other
6//! length by Bluestein's chirp-z algorithm, which expresses the DFT as a
7//! convolution and evaluates that with a power-of-two FFT of size ≥ 2n − 1.
8//! Both are O(n log n); the naive DFT is O(n²).
9//!
10//! Conventions: the forward transform is `X_k = Σ_j x_j · exp(−2πi jk/n)`
11//! (no scaling) and the inverse divides by `n`, so `ifft(fft(x)) == x`.
12//!
13//! # References
14//! - Cooley, J. W. & Tukey, J. W. (1965). "An algorithm for the machine
15//!   calculation of complex Fourier series." *Mathematics of Computation*,
16//!   19(90), 297–301.
17//! - Bluestein, L. I. (1970). "A linear filtering approach to the computation
18//!   of discrete Fourier transform." *IEEE Transactions on Audio and
19//!   Electroacoustics*, 18(4), 451–455.
20
21use std::f64::consts::PI;
22use std::ops::{Add, Mul, Sub};
23
24/// A complex number with `f64` parts.
25#[derive(Debug, Clone, Copy, PartialEq, Default)]
26pub struct Complex {
27    /// Real part.
28    pub re: f64,
29    /// Imaginary part.
30    pub im: f64,
31}
32
33impl Complex {
34    /// Creates `re + im·i`.
35    pub const fn new(re: f64, im: f64) -> Self {
36        Self { re, im }
37    }
38
39    /// Creates a real number (imaginary part 0).
40    pub const fn real(re: f64) -> Self {
41        Self { re, im: 0.0 }
42    }
43
44    /// `exp(i·theta)`.
45    pub fn from_polar_unit(theta: f64) -> Self {
46        Self {
47            re: theta.cos(),
48            im: theta.sin(),
49        }
50    }
51
52    /// Complex conjugate.
53    pub fn conj(self) -> Self {
54        Self {
55            re: self.re,
56            im: -self.im,
57        }
58    }
59
60    /// Modulus `|z|`.
61    pub fn norm(self) -> f64 {
62        self.re.hypot(self.im)
63    }
64
65    /// Squared modulus `|z|²`.
66    pub fn norm_sqr(self) -> f64 {
67        self.re * self.re + self.im * self.im
68    }
69
70    /// Argument in radians, in `(−π, π]`.
71    pub fn arg(self) -> f64 {
72        self.im.atan2(self.re)
73    }
74
75    /// Multiplies by a real scalar.
76    pub fn scale(self, k: f64) -> Self {
77        Self {
78            re: self.re * k,
79            im: self.im * k,
80        }
81    }
82}
83
84impl Add for Complex {
85    type Output = Self;
86    fn add(self, rhs: Self) -> Self {
87        Self {
88            re: self.re + rhs.re,
89            im: self.im + rhs.im,
90        }
91    }
92}
93
94impl Sub for Complex {
95    type Output = Self;
96    fn sub(self, rhs: Self) -> Self {
97        Self {
98            re: self.re - rhs.re,
99            im: self.im - rhs.im,
100        }
101    }
102}
103
104impl Mul for Complex {
105    type Output = Self;
106    fn mul(self, rhs: Self) -> Self {
107        Self {
108            re: self.re * rhs.re - self.im * rhs.im,
109            im: self.re * rhs.im + self.im * rhs.re,
110        }
111    }
112}
113
114/// Forward DFT of a complex sequence of any length.
115///
116/// Returns `X_k = Σ_j x_j · exp(−2πi jk/n)` for `k = 0..n`. An empty input
117/// returns an empty output.
118///
119/// # Complexity
120/// O(n log n) — radix-2 for a power-of-two `n`, Bluestein otherwise.
121///
122/// # Examples
123///
124/// ```
125/// use u_numflow::fourier::{fft, Complex};
126///
127/// // A constant sequence has all its energy in the zero frequency.
128/// let x = vec![Complex::real(1.0); 4];
129/// let spectrum = fft(&x);
130/// assert!((spectrum[0].re - 4.0).abs() < 1e-12);
131/// assert!(spectrum[1..].iter().all(|z| z.norm() < 1e-12));
132/// ```
133pub fn fft(input: &[Complex]) -> Vec<Complex> {
134    transform(input, false)
135}
136
137/// Inverse DFT: `x_j = (1/n) Σ_k X_k · exp(+2πi jk/n)`.
138///
139/// `ifft(fft(x))` returns `x` up to floating-point rounding.
140///
141/// # Examples
142///
143/// ```
144/// use u_numflow::fourier::{fft, ifft, Complex};
145///
146/// let x: Vec<Complex> = (0..5).map(|j| Complex::new(j as f64, -(j as f64))).collect();
147/// let back = ifft(&fft(&x));
148/// for (a, b) in x.iter().zip(&back) {
149///     assert!((a.re - b.re).abs() < 1e-12 && (a.im - b.im).abs() < 1e-12);
150/// }
151/// ```
152pub fn ifft(input: &[Complex]) -> Vec<Complex> {
153    let n = input.len();
154    if n == 0 {
155        return Vec::new();
156    }
157    let scale = 1.0 / n as f64;
158    transform(input, true)
159        .into_iter()
160        .map(|z| z.scale(scale))
161        .collect()
162}
163
164/// Forward DFT of a real sequence: the `n / 2 + 1` non-redundant bins
165/// `0..=n/2` (NumPy `numpy.fft.rfft`, SciPy `scipy.fft.rfft`, FFTW r2c).
166///
167/// For real input the rest mirrors them (`X_{n−k} = conj(X_k)`), so nothing
168/// is lost; [`fft`] of the same values, taken as complex, returns all `n`.
169/// An empty input gives an empty spectrum.
170///
171/// # Examples
172///
173/// ```
174/// use u_numflow::fourier::rfft;
175///
176/// // One cycle of a cosine over 8 samples: all its energy is in bin 1.
177/// let x: Vec<f64> = (0..8).map(|j| (2.0 * std::f64::consts::PI * j as f64 / 8.0).cos()).collect();
178/// let spectrum = rfft(&x);
179/// assert_eq!(spectrum.len(), 5);
180/// assert!((spectrum[1].re - 4.0).abs() < 1e-12);
181/// ```
182pub fn rfft(input: &[f64]) -> Vec<Complex> {
183    let complex: Vec<Complex> = input.iter().map(|&x| Complex::real(x)).collect();
184    let mut spectrum = fft(&complex);
185    spectrum.truncate(input.len() / 2 + 1);
186    spectrum
187}
188
189/// Dispatches on the length: radix-2 for powers of two, Bluestein otherwise.
190fn transform(input: &[Complex], inverse: bool) -> Vec<Complex> {
191    let n = input.len();
192    if n <= 1 {
193        return input.to_vec();
194    }
195    let mut data = input.to_vec();
196    if n.is_power_of_two() {
197        radix2_in_place(&mut data, inverse);
198        data
199    } else {
200        bluestein(&data, inverse)
201    }
202}
203
204/// Iterative radix-2 Cooley-Tukey; `data.len()` must be a power of two.
205fn radix2_in_place(data: &mut [Complex], inverse: bool) {
206    let n = data.len();
207    debug_assert!(n.is_power_of_two());
208
209    // Bit-reversal permutation.
210    let bits = n.trailing_zeros();
211    for i in 0..n {
212        let j = i.reverse_bits() >> (usize::BITS - bits);
213        if j > i {
214            data.swap(i, j);
215        }
216    }
217
218    let sign = if inverse { 1.0 } else { -1.0 };
219    let mut len = 2;
220    while len <= n {
221        let step = Complex::from_polar_unit(sign * 2.0 * PI / len as f64);
222        for start in (0..n).step_by(len) {
223            let mut w = Complex::real(1.0);
224            for k in 0..len / 2 {
225                let a = data[start + k];
226                let b = data[start + k + len / 2] * w;
227                data[start + k] = a + b;
228                data[start + k + len / 2] = a - b;
229                w = w * step;
230            }
231        }
232        len <<= 1;
233    }
234}
235
236/// Bluestein's chirp-z algorithm for an arbitrary length `n`.
237///
238/// With the chirp `w_j = exp(∓πi j²/n)`, `X_k = w_k · Σ_j (x_j w_j) · conj(w_{k−j})`
239/// is a linear convolution, evaluated as a circular one of power-of-two size
240/// `m ≥ 2n − 1`.
241fn bluestein(input: &[Complex], inverse: bool) -> Vec<Complex> {
242    let n = input.len();
243    let m = (2 * n - 1).next_power_of_two();
244    let sign = if inverse { 1.0 } else { -1.0 };
245
246    // Chirp: the angle is π·j²/n; j² is reduced modulo 2n first so the
247    // argument stays exact for large j.
248    let chirp: Vec<Complex> = (0..n)
249        .map(|j| {
250            let jj = (j * j) % (2 * n);
251            Complex::from_polar_unit(sign * PI * jj as f64 / n as f64)
252        })
253        .collect();
254
255    let mut a = vec![Complex::default(); m];
256    for j in 0..n {
257        a[j] = input[j] * chirp[j];
258    }
259    let mut b = vec![Complex::default(); m];
260    b[0] = chirp[0].conj();
261    for j in 1..n {
262        let c = chirp[j].conj();
263        b[j] = c;
264        b[m - j] = c;
265    }
266
267    radix2_in_place(&mut a, false);
268    radix2_in_place(&mut b, false);
269    let mut c: Vec<Complex> = a.iter().zip(&b).map(|(&x, &y)| x * y).collect();
270    radix2_in_place(&mut c, true);
271    let scale = 1.0 / m as f64;
272
273    (0..n).map(|k| c[k].scale(scale) * chirp[k]).collect()
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279    use proptest::prelude::*;
280
281    fn naive_dft(x: &[Complex], inverse: bool) -> Vec<Complex> {
282        let n = x.len();
283        let sign = if inverse { 1.0 } else { -1.0 };
284        (0..n)
285            .map(|k| {
286                let mut acc = Complex::default();
287                for (j, &xj) in x.iter().enumerate() {
288                    let angle = sign * 2.0 * PI * (j * k) as f64 / n as f64;
289                    acc = acc + xj * Complex::from_polar_unit(angle);
290                }
291                if inverse {
292                    acc.scale(1.0 / n as f64)
293                } else {
294                    acc
295                }
296            })
297            .collect()
298    }
299
300    fn assert_close(a: &[Complex], b: &[Complex], tol: f64) {
301        assert_eq!(a.len(), b.len());
302        for (i, (x, y)) in a.iter().zip(b).enumerate() {
303            assert!(
304                (*x - *y).norm() < tol,
305                "bin {i}: {x:?} vs {y:?} (|Δ| = {})",
306                (*x - *y).norm()
307            );
308        }
309    }
310
311    #[test]
312    fn empty_and_singleton() {
313        assert!(fft(&[]).is_empty());
314        assert!(ifft(&[]).is_empty());
315        let one = [Complex::new(2.0, -3.0)];
316        assert_eq!(fft(&one), one.to_vec());
317        assert_eq!(ifft(&one), one.to_vec());
318    }
319
320    #[test]
321    fn impulse_is_flat() {
322        let mut x = vec![Complex::default(); 7];
323        x[0] = Complex::real(1.0);
324        for z in fft(&x) {
325            assert!((z.re - 1.0).abs() < 1e-12 && z.im.abs() < 1e-12);
326        }
327    }
328
329    #[test]
330    fn sine_lands_in_its_bin_for_every_length_kind() {
331        for n in [8usize, 12, 30, 64, 97] {
332            let x: Vec<f64> = (0..n)
333                .map(|j| (2.0 * PI * 3.0 * j as f64 / n as f64).sin())
334                .collect();
335            let spectrum = rfft(&x);
336            let power: Vec<f64> = spectrum.iter().map(|z| z.norm_sqr()).collect();
337            let (argmax, _) = power[..n / 2 + 1]
338                .iter()
339                .enumerate()
340                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
341                .unwrap();
342            assert_eq!(argmax, 3, "n = {n}");
343            assert!((spectrum[3].im + n as f64 / 2.0).abs() < 1e-9, "n = {n}");
344        }
345    }
346
347    proptest! {
348        #![proptest_config(ProptestConfig::with_cases(128))]
349
350        #[test]
351        fn matches_the_naive_dft_for_any_length(
352            data in prop::collection::vec((-100.0f64..100.0, -100.0f64..100.0), 1..=64)
353        ) {
354            let x: Vec<Complex> = data.iter().map(|&(re, im)| Complex::new(re, im)).collect();
355            let n = x.len() as f64;
356            let tol = 1e-9 * n * 100.0;
357            let forward = fft(&x);
358            assert_close(&forward, &naive_dft(&x, false), tol);
359            assert_close(&ifft(&forward), &x, tol);
360            assert_close(&ifft(&x), &naive_dft(&x, true), tol);
361        }
362
363        /// Parseval: Σ|x_j|² = (1/n) Σ|X_k|².
364        #[test]
365        fn parseval_holds(
366            data in prop::collection::vec(-100.0f64..100.0, 1..=100)
367        ) {
368            let n = data.len() as f64;
369            let time: f64 = data.iter().map(|x| x * x).sum();
370            // The full spectrum, rebuilt from the half rfft returns: every bin
371            // between DC and Nyquist stands for itself and its mirror.
372            let half = rfft(&data);
373            let len = data.len();
374            let freq: f64 = half
375                .iter()
376                .enumerate()
377                .map(|(k, z)| {
378                    let twice = k != 0 && 2 * k != len;
379                    z.norm_sqr() * if twice { 2.0 } else { 1.0 }
380                })
381                .sum::<f64>()
382                / n;
383            prop_assert!((time - freq).abs() < 1e-7 * (1.0 + time), "{time} vs {freq}");
384        }
385
386        /// Real input: `rfft` is the first `n / 2 + 1` bins of `fft`, and the
387        /// bins it leaves out are their conjugate mirror.
388        #[test]
389        fn real_input_has_hermitian_spectrum(
390            data in prop::collection::vec(-100.0f64..100.0, 2..=100)
391        ) {
392            let n = data.len();
393            let half = rfft(&data);
394            prop_assert_eq!(half.len(), n / 2 + 1);
395            let full = fft(&data.iter().map(|&x| Complex::real(x)).collect::<Vec<_>>());
396            for k in 0..half.len() {
397                prop_assert!((half[k] - full[k]).norm() < 1e-9 * (1.0 + full[k].norm()));
398            }
399            for k in 1..n {
400                let mirror = full[n - k].conj();
401                prop_assert!((full[k] - mirror).norm() < 1e-7 * (1.0 + full[k].norm()));
402            }
403            prop_assert!(half[0].im.abs() < 1e-9);
404        }
405    }
406}
407
408/// Against NumPy 2.4.6 (`numpy.fft.rfft`).
409#[cfg(test)]
410mod numpy_reference {
411    use super::*;
412
413    #[test]
414    fn rfft_of_an_odd_length_reproduces_numpy() {
415        // numpy.fft.rfft([0.5, 1.5, -0.25, 2.0, 0.0, -1.0, 3.5, 1.25, -0.75]) -- length 9,
416        // not a power of two. Both return the bins 0..=n/2.
417        let got = rfft(&[0.5, 1.5, -0.25, 2.0, 0.0, -1.0, 3.5, 1.25, -0.75]);
418        let want = [
419            (6.75, 0.0),
420            (-0.5621258692079278, 0.9879574700745879),
421            (-3.8255009306546883, -2.3590377252790837),
422            (5.625, -4.113620667976083),
423            (-2.362373200137384, 0.5501191216763025),
424        ];
425        assert_eq!(got.len(), 5);
426        for (k, (re, im)) in want.into_iter().enumerate() {
427            let g = got[k];
428            assert!(
429                (g.re - re).abs() <= 1e-12 && (g.im - im).abs() <= 1e-12,
430                "bin {k}: got {} + {}i, NumPy {re} + {im}i",
431                g.re,
432                g.im
433            );
434        }
435    }
436}
u_numflow/fourier.rs

u_numflow/
fourier.rs
