Skip to main content

oxiphysics_core/
signal_processing.rs

1#![allow(clippy::needless_range_loop)]
2// Copyright 2026 COOLJAPAN OU (Team KitaSan)
3// SPDX-License-Identifier: Apache-2.0
4
5//! Signal processing: FFT, windowing, filtering, spectral analysis, wavelets.
6//!
7//! Provides Cooley-Tukey FFT/IFFT, windowing functions (Hann/Hamming/Blackman/Kaiser),
8//! PSD estimation (Welch/Bartlett), FIR/IIR filter design (Butterworth/Chebyshev/elliptic),
9//! convolution, correlation, Hilbert transform, spectrogram, and wavelet transforms.
10
11#![allow(dead_code)]
12#![allow(unused_imports)]
13#![allow(clippy::too_many_arguments)]
14
15use std::f64::consts::PI;
16
17// ---------------------------------------------------------------------------
18// Complex number (local, minimal)
19// ---------------------------------------------------------------------------
20
21/// A complex number with `f64` real and imaginary parts.
22#[derive(Clone, Copy, Debug, PartialEq)]
23pub struct Complex64 {
24    /// Real part.
25    pub re: f64,
26    /// Imaginary part.
27    pub im: f64,
28}
29
30impl Complex64 {
31    /// Construct a new `Complex64`.
32    #[inline]
33    pub fn new(re: f64, im: f64) -> Self {
34        Self { re, im }
35    }
36
37    /// Construct from magnitude and phase (polar form).
38    #[inline]
39    pub fn from_polar(r: f64, theta: f64) -> Self {
40        Self {
41            re: r * theta.cos(),
42            im: r * theta.sin(),
43        }
44    }
45
46    /// Return the complex conjugate.
47    #[inline]
48    pub fn conj(self) -> Self {
49        Self {
50            re: self.re,
51            im: -self.im,
52        }
53    }
54
55    /// Return the modulus (absolute value).
56    #[inline]
57    pub fn abs(self) -> f64 {
58        self.re.hypot(self.im)
59    }
60
61    /// Return the argument (phase angle in radians).
62    #[inline]
63    pub fn arg(self) -> f64 {
64        self.im.atan2(self.re)
65    }
66
67    /// Return the squared modulus.
68    #[inline]
69    pub fn norm_sq(self) -> f64 {
70        self.re * self.re + self.im * self.im
71    }
72
73    /// Multiply two complex numbers.
74    #[inline]
75    #[allow(clippy::should_implement_trait)]
76    pub fn mul(self, rhs: Self) -> Self {
77        Self {
78            re: self.re * rhs.re - self.im * rhs.im,
79            im: self.re * rhs.im + self.im * rhs.re,
80        }
81    }
82
83    /// Add two complex numbers.
84    #[inline]
85    #[allow(clippy::should_implement_trait)]
86    pub fn add(self, rhs: Self) -> Self {
87        Self {
88            re: self.re + rhs.re,
89            im: self.im + rhs.im,
90        }
91    }
92
93    /// Subtract two complex numbers.
94    #[inline]
95    #[allow(clippy::should_implement_trait)]
96    pub fn sub(self, rhs: Self) -> Self {
97        Self {
98            re: self.re - rhs.re,
99            im: self.im - rhs.im,
100        }
101    }
102
103    /// Scale by a real scalar.
104    #[inline]
105    pub fn scale(self, s: f64) -> Self {
106        Self {
107            re: self.re * s,
108            im: self.im * s,
109        }
110    }
111
112    /// Compute e^(i*theta) = cos(theta) + i*sin(theta).
113    #[inline]
114    pub fn exp_i(theta: f64) -> Self {
115        Self {
116            re: theta.cos(),
117            im: theta.sin(),
118        }
119    }
120
121    /// Compute the natural exponential e^z.
122    #[inline]
123    pub fn exp(self) -> Self {
124        let r = self.re.exp();
125        Self {
126            re: r * self.im.cos(),
127            im: r * self.im.sin(),
128        }
129    }
130}
131
132impl std::ops::Add for Complex64 {
133    type Output = Self;
134    fn add(self, rhs: Self) -> Self {
135        Self::add(self, rhs)
136    }
137}
138
139impl std::ops::Sub for Complex64 {
140    type Output = Self;
141    fn sub(self, rhs: Self) -> Self {
142        Self::sub(self, rhs)
143    }
144}
145
146impl std::ops::Mul for Complex64 {
147    type Output = Self;
148    fn mul(self, rhs: Self) -> Self {
149        Self::mul(self, rhs)
150    }
151}
152
153impl std::ops::Neg for Complex64 {
154    type Output = Self;
155    fn neg(self) -> Self {
156        Self {
157            re: -self.re,
158            im: -self.im,
159        }
160    }
161}
162
163// ---------------------------------------------------------------------------
164// FFT / IFFT — Cooley-Tukey radix-2
165// ---------------------------------------------------------------------------
166
167/// Bit-reverse the indices of a slice in place.
168fn bit_reverse_copy(a: &mut [Complex64]) {
169    let n = a.len();
170    let mut j = 0usize;
171    for i in 1..n {
172        let mut bit = n >> 1;
173        while j & bit != 0 {
174            j ^= bit;
175            bit >>= 1;
176        }
177        j ^= bit;
178        if i < j {
179            a.swap(i, j);
180        }
181    }
182}
183
184/// In-place Cooley-Tukey FFT.  `n` must be a power of two.
185///
186/// Uses the DIT (decimation-in-time) butterfly algorithm.
187pub fn fft_inplace(a: &mut [Complex64]) {
188    let n = a.len();
189    assert!(n.is_power_of_two(), "FFT length must be a power of two");
190    bit_reverse_copy(a);
191    let mut len = 2usize;
192    while len <= n {
193        let ang = -2.0 * PI / len as f64;
194        let wlen = Complex64::exp_i(ang);
195        let mut i = 0;
196        while i < n {
197            let mut w = Complex64::new(1.0, 0.0);
198            for j in 0..(len / 2) {
199                let u = a[i + j];
200                let v = a[i + j + len / 2] * w;
201                a[i + j] = u + v;
202                a[i + j + len / 2] = u - v;
203                w = w * wlen;
204            }
205            i += len;
206        }
207        len <<= 1;
208    }
209}
210
211/// In-place inverse FFT (IFFT).
212pub fn ifft_inplace(a: &mut [Complex64]) {
213    let n = a.len();
214    // conjugate
215    for x in a.iter_mut() {
216        x.im = -x.im;
217    }
218    fft_inplace(a);
219    // conjugate and scale
220    let inv = 1.0 / n as f64;
221    for x in a.iter_mut() {
222        x.im = -x.im;
223        x.re *= inv;
224        x.im *= inv;
225    }
226}
227
228/// Compute FFT of a real-valued signal.  Pads to the next power of two if needed.
229///
230/// Returns the full complex spectrum of length equal to the next power of two.
231pub fn fft(signal: &[f64]) -> Vec<Complex64> {
232    let n = signal.len().next_power_of_two();
233    let mut buf: Vec<Complex64> = signal
234        .iter()
235        .map(|&x| Complex64::new(x, 0.0))
236        .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
237        .take(n)
238        .collect();
239    fft_inplace(&mut buf);
240    buf
241}
242
243/// Compute IFFT and return the real part of each element.
244pub fn ifft_real(spectrum: &[Complex64]) -> Vec<f64> {
245    let mut buf = spectrum.to_vec();
246    ifft_inplace(&mut buf);
247    buf.iter().map(|c| c.re).collect()
248}
249
250/// Compute the power spectrum (|X\[k\]|^2) from a real signal.
251pub fn power_spectrum(signal: &[f64]) -> Vec<f64> {
252    let spec = fft(signal);
253    spec.iter().map(|c| c.norm_sq()).collect()
254}
255
256/// Compute the magnitude spectrum (|X\[k\]|) from a real signal.
257pub fn magnitude_spectrum(signal: &[f64]) -> Vec<f64> {
258    let spec = fft(signal);
259    spec.iter().map(|c| c.abs()).collect()
260}
261
262/// Compute the phase spectrum (arg(X\[k\])) from a real signal.
263pub fn phase_spectrum(signal: &[f64]) -> Vec<f64> {
264    let spec = fft(signal);
265    spec.iter().map(|c| c.arg()).collect()
266}
267
268/// Return the frequency bins (in Hz) for an FFT of length `n` sampled at `fs` Hz.
269pub fn fft_frequencies(n: usize, fs: f64) -> Vec<f64> {
270    (0..n).map(|k| k as f64 * fs / n as f64).collect()
271}
272
273/// Return only the positive-frequency bins (0 to fs/2).
274pub fn rfft_frequencies(n: usize, fs: f64) -> Vec<f64> {
275    (0..=n / 2).map(|k| k as f64 * fs / n as f64).collect()
276}
277
278// ---------------------------------------------------------------------------
279// Windowing functions
280// ---------------------------------------------------------------------------
281
282/// The type of window function to apply.
283#[derive(Clone, Copy, Debug, PartialEq)]
284pub enum WindowType {
285    /// Rectangular (no windowing).
286    Rectangular,
287    /// Hann window.
288    Hann,
289    /// Hamming window.
290    Hamming,
291    /// Blackman window.
292    Blackman,
293    /// Blackman-Harris window.
294    BlackmanHarris,
295    /// Flat-top window.
296    FlatTop,
297    /// Bartlett (triangular) window.
298    Bartlett,
299    /// Kaiser window (parameter beta).
300    Kaiser,
301    /// Nuttall window.
302    Nuttall,
303}
304
305/// Compute the Hann window of length `n`.
306pub fn hann_window(n: usize) -> Vec<f64> {
307    (0..n)
308        .map(|i| 0.5 * (1.0 - (2.0 * PI * i as f64 / (n - 1) as f64).cos()))
309        .collect()
310}
311
312/// Compute the Hamming window of length `n`.
313pub fn hamming_window(n: usize) -> Vec<f64> {
314    (0..n)
315        .map(|i| 0.54 - 0.46 * (2.0 * PI * i as f64 / (n - 1) as f64).cos())
316        .collect()
317}
318
319/// Compute the Blackman window of length `n`.
320pub fn blackman_window(n: usize) -> Vec<f64> {
321    (0..n)
322        .map(|i| {
323            let t = 2.0 * PI * i as f64 / (n - 1) as f64;
324            0.42 - 0.5 * t.cos() + 0.08 * (2.0 * t).cos()
325        })
326        .collect()
327}
328
329/// Compute the Blackman-Harris window of length `n`.
330pub fn blackman_harris_window(n: usize) -> Vec<f64> {
331    let a0 = 0.35875;
332    let a1 = 0.48829;
333    let a2 = 0.14128;
334    let a3 = 0.01168;
335    (0..n)
336        .map(|i| {
337            let t = 2.0 * PI * i as f64 / (n - 1) as f64;
338            a0 - a1 * t.cos() + a2 * (2.0 * t).cos() - a3 * (3.0 * t).cos()
339        })
340        .collect()
341}
342
343/// Compute the Nuttall window of length `n`.
344pub fn nuttall_window(n: usize) -> Vec<f64> {
345    let a0 = 0.3635819;
346    let a1 = 0.4891775;
347    let a2 = 0.1365995;
348    let a3 = 0.0106411;
349    (0..n)
350        .map(|i| {
351            let t = 2.0 * PI * i as f64 / (n - 1) as f64;
352            a0 - a1 * t.cos() + a2 * (2.0 * t).cos() - a3 * (3.0 * t).cos()
353        })
354        .collect()
355}
356
357/// Compute the Bartlett (triangular) window of length `n`.
358pub fn bartlett_window(n: usize) -> Vec<f64> {
359    let m = (n - 1) as f64;
360    (0..n)
361        .map(|i| {
362            let x = i as f64;
363            1.0 - (2.0 * x / m - 1.0).abs()
364        })
365        .collect()
366}
367
368/// Compute the flat-top window of length `n`.
369pub fn flat_top_window(n: usize) -> Vec<f64> {
370    let a0 = 0.21557895;
371    let a1 = 0.41663158;
372    let a2 = 0.277263158;
373    let a3 = 0.083578947;
374    let a4 = 0.006947368;
375    (0..n)
376        .map(|i| {
377            let t = 2.0 * PI * i as f64 / (n - 1) as f64;
378            a0 - a1 * t.cos() + a2 * (2.0 * t).cos() - a3 * (3.0 * t).cos() + a4 * (4.0 * t).cos()
379        })
380        .collect()
381}
382
383/// Modified Bessel function of the first kind, order zero (I₀).
384/// Used for Kaiser window computation.
385pub fn bessel_i0(x: f64) -> f64 {
386    // Series expansion: I0(x) = sum_{k=0}^{inf} (x/2)^{2k} / (k!)^2
387    let mut sum = 1.0;
388    let mut term = 1.0;
389    let half_x = x / 2.0;
390    for k in 1..=50 {
391        term *= half_x * half_x / (k * k) as f64;
392        sum += term;
393        if term < 1e-15 * sum {
394            break;
395        }
396    }
397    sum
398}
399
400/// Compute the Kaiser window of length `n` with shape parameter `beta`.
401pub fn kaiser_window(n: usize, beta: f64) -> Vec<f64> {
402    let i0_beta = bessel_i0(beta);
403    let m = (n - 1) as f64;
404    (0..n)
405        .map(|i| {
406            let t = 2.0 * i as f64 / m - 1.0;
407            let arg = beta * (1.0 - t * t).max(0.0).sqrt();
408            bessel_i0(arg) / i0_beta
409        })
410        .collect()
411}
412
413/// Apply a window to a signal in-place.
414pub fn apply_window(signal: &mut [f64], window: &[f64]) {
415    assert_eq!(
416        signal.len(),
417        window.len(),
418        "Signal and window must have equal length"
419    );
420    for (s, w) in signal.iter_mut().zip(window.iter()) {
421        *s *= w;
422    }
423}
424
425/// Apply a window and return the windowed signal.
426pub fn windowed(signal: &[f64], window: &[f64]) -> Vec<f64> {
427    assert_eq!(signal.len(), window.len());
428    signal
429        .iter()
430        .zip(window.iter())
431        .map(|(s, w)| s * w)
432        .collect()
433}
434
435/// Compute the coherent gain of a window (mean).
436pub fn window_coherent_gain(window: &[f64]) -> f64 {
437    window.iter().sum::<f64>() / window.len() as f64
438}
439
440/// Compute the power gain of a window (mean of squares).
441pub fn window_power_gain(window: &[f64]) -> f64 {
442    window.iter().map(|w| w * w).sum::<f64>() / window.len() as f64
443}
444
445// ---------------------------------------------------------------------------
446// Power Spectral Density estimation
447// ---------------------------------------------------------------------------
448
449/// PSD estimation result containing frequency bins and power values.
450#[derive(Clone, Debug)]
451pub struct PsdResult {
452    /// Frequency values (Hz).
453    pub freqs: Vec<f64>,
454    /// Power spectral density (units²/Hz).
455    pub psd: Vec<f64>,
456}
457
458/// Welch's method for PSD estimation.
459///
460/// Splits the signal into overlapping segments, windows each, computes the
461/// periodogram, and averages.
462pub fn welch_psd(
463    signal: &[f64],
464    fs: f64,
465    segment_len: usize,
466    overlap: usize,
467    window_type: WindowType,
468) -> PsdResult {
469    let step = segment_len - overlap;
470    let window = match window_type {
471        WindowType::Hann => hann_window(segment_len),
472        WindowType::Hamming => hamming_window(segment_len),
473        WindowType::Blackman => blackman_window(segment_len),
474        WindowType::Bartlett => bartlett_window(segment_len),
475        _ => hann_window(segment_len),
476    };
477    let win_power: f64 = window.iter().map(|w| w * w).sum::<f64>();
478
479    let n_fft = segment_len.next_power_of_two();
480    let n_bins = n_fft / 2 + 1;
481    let mut psd_acc = vec![0.0f64; n_bins];
482    let mut seg_count = 0usize;
483
484    let mut start = 0;
485    while start + segment_len <= signal.len() {
486        let seg = &signal[start..start + segment_len];
487        let mut windowed_seg: Vec<f64> =
488            seg.iter().zip(window.iter()).map(|(s, w)| s * w).collect();
489        windowed_seg.resize(n_fft, 0.0);
490        let spec = fft(&windowed_seg);
491        // Periodogram: |X[k]|^2 / (fs * W)
492        for k in 0..n_bins {
493            let factor = if k == 0 || k == n_fft / 2 { 1.0 } else { 2.0 };
494            psd_acc[k] += factor * spec[k].norm_sq() / (fs * win_power);
495        }
496        seg_count += 1;
497        start += step;
498    }
499    // Average over segments
500    if seg_count > 0 {
501        let inv = 1.0 / seg_count as f64;
502        for p in psd_acc.iter_mut() {
503            *p *= inv;
504        }
505    }
506    let freqs = rfft_frequencies(n_fft, fs);
507    PsdResult {
508        freqs,
509        psd: psd_acc,
510    }
511}
512
513/// Bartlett's method for PSD estimation (non-overlapping segments).
514pub fn bartlett_psd(signal: &[f64], fs: f64, segment_len: usize) -> PsdResult {
515    welch_psd(signal, fs, segment_len, 0, WindowType::Rectangular)
516}
517
518/// Simple periodogram PSD estimate.
519pub fn periodogram_psd(signal: &[f64], fs: f64, window_type: WindowType) -> PsdResult {
520    welch_psd(signal, fs, signal.len(), 0, window_type)
521}
522
523// ---------------------------------------------------------------------------
524// Convolution and Correlation
525// ---------------------------------------------------------------------------
526
527/// Linear convolution of two real signals via FFT (overlap-add style).
528pub fn convolve(x: &[f64], h: &[f64]) -> Vec<f64> {
529    let out_len = x.len() + h.len() - 1;
530    let n = out_len.next_power_of_two();
531    let mut xc: Vec<Complex64> = x
532        .iter()
533        .map(|&v| Complex64::new(v, 0.0))
534        .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
535        .take(n)
536        .collect();
537    let mut hc: Vec<Complex64> = h
538        .iter()
539        .map(|&v| Complex64::new(v, 0.0))
540        .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
541        .take(n)
542        .collect();
543    fft_inplace(&mut xc);
544    fft_inplace(&mut hc);
545    let mut yc: Vec<Complex64> = xc.iter().zip(hc.iter()).map(|(a, b)| *a * *b).collect();
546    ifft_inplace(&mut yc);
547    yc[..out_len].iter().map(|c| c.re).collect()
548}
549
550/// Cross-correlation of `x` with `y` (i.e. sum_t x\[t\] * y\[t+lag\]).
551pub fn cross_correlate(x: &[f64], y: &[f64]) -> Vec<f64> {
552    // Flip y and convolve
553    let y_flip: Vec<f64> = y.iter().rev().cloned().collect();
554    convolve(x, &y_flip)
555}
556
557/// Auto-correlation of `x`.
558pub fn auto_correlate(x: &[f64]) -> Vec<f64> {
559    cross_correlate(x, x)
560}
561
562/// Normalized cross-correlation coefficient at each lag.
563pub fn normalized_cross_correlation(x: &[f64], y: &[f64]) -> Vec<f64> {
564    let xcorr = cross_correlate(x, y);
565    let norm = (x.iter().map(|v| v * v).sum::<f64>() * y.iter().map(|v| v * v).sum::<f64>()).sqrt();
566    if norm == 0.0 {
567        xcorr
568    } else {
569        xcorr.iter().map(|v| v / norm).collect()
570    }
571}
572
573// ---------------------------------------------------------------------------
574// FIR Filter Design
575// ---------------------------------------------------------------------------
576
577/// FIR filter coefficients computed via the window method.
578#[derive(Clone, Debug)]
579pub struct FirFilter {
580    /// Filter taps (coefficients).
581    pub coeffs: Vec<f64>,
582}
583
584impl FirFilter {
585    /// Design a low-pass FIR filter using the window method.
586    ///
587    /// `n_taps`: number of taps (should be odd for linear phase).
588    /// `cutoff`: normalised cutoff frequency (0..1, where 1 = Nyquist).
589    pub fn lowpass(n_taps: usize, cutoff: f64) -> Self {
590        let n_taps = if n_taps.is_multiple_of(2) {
591            n_taps + 1
592        } else {
593            n_taps
594        };
595        let m = (n_taps - 1) as f64;
596        let fc = cutoff * PI;
597        let window = hann_window(n_taps);
598        let coeffs: Vec<f64> = (0..n_taps)
599            .map(|i| {
600                let n = i as f64 - m / 2.0;
601                if n == 0.0 {
602                    fc / PI
603                } else {
604                    (fc * n).sin() / (PI * n) * window[i]
605                }
606            })
607            .collect();
608        Self { coeffs }
609    }
610
611    /// Design a high-pass FIR filter.
612    pub fn highpass(n_taps: usize, cutoff: f64) -> Self {
613        let mut lp = Self::lowpass(n_taps, cutoff);
614        // Spectral inversion
615        for (i, c) in lp.coeffs.iter_mut().enumerate() {
616            *c = if i == (n_taps - 1) / 2 { 1.0 - *c } else { -*c };
617        }
618        lp
619    }
620
621    /// Design a band-pass FIR filter.
622    pub fn bandpass(n_taps: usize, low: f64, high: f64) -> Self {
623        let lp_high = Self::lowpass(n_taps, high);
624        let lp_low = Self::lowpass(n_taps, low);
625        let coeffs: Vec<f64> = lp_high
626            .coeffs
627            .iter()
628            .zip(lp_low.coeffs.iter())
629            .map(|(h, l)| h - l)
630            .collect();
631        Self { coeffs }
632    }
633
634    /// Design a band-stop (notch) FIR filter.
635    pub fn bandstop(n_taps: usize, low: f64, high: f64) -> Self {
636        let bp = Self::bandpass(n_taps, low, high);
637        let n_taps = bp.coeffs.len();
638        let coeffs: Vec<f64> = bp
639            .coeffs
640            .iter()
641            .enumerate()
642            .map(|(i, c)| if i == (n_taps - 1) / 2 { 1.0 - c } else { -*c })
643            .collect();
644        Self { coeffs }
645    }
646
647    /// Apply the FIR filter to a signal (direct-form convolution).
648    pub fn apply(&self, signal: &[f64]) -> Vec<f64> {
649        convolve(signal, &self.coeffs)
650    }
651
652    /// Apply using overlap-save direct implementation (causal, same length as input).
653    pub fn filter(&self, signal: &[f64]) -> Vec<f64> {
654        let n = self.coeffs.len();
655        let out_full = self.apply(signal);
656        let delay = (n - 1) / 2;
657        // Return signal aligned (no group delay)
658        if out_full.len() >= signal.len() + delay {
659            out_full[delay..delay + signal.len()].to_vec()
660        } else {
661            let end = out_full.len().min(signal.len());
662            out_full[..end].to_vec()
663        }
664    }
665
666    /// Compute the frequency response at `n_points` equally spaced frequencies \[0, fs/2\].
667    pub fn frequency_response(&self, n_points: usize) -> Vec<Complex64> {
668        let n_fft = n_points.next_power_of_two() * 2;
669        let mut buf: Vec<Complex64> = self
670            .coeffs
671            .iter()
672            .map(|&c| Complex64::new(c, 0.0))
673            .chain(std::iter::repeat(Complex64::new(0.0, 0.0)))
674            .take(n_fft)
675            .collect();
676        fft_inplace(&mut buf);
677        buf[..n_points].to_vec()
678    }
679}
680
681// ---------------------------------------------------------------------------
682// IIR Filter — Biquad sections
683// ---------------------------------------------------------------------------
684
685/// A single biquad (second-order) IIR section.
686#[derive(Clone, Debug)]
687pub struct Biquad {
688    /// Feed-forward coefficients: b0, b1, b2.
689    pub b: [f64; 3],
690    /// Feedback coefficients: a1, a2 (a0 = 1 normalised).
691    pub a: [f64; 2],
692    /// Filter state (delay line).
693    state: [f64; 2],
694}
695
696impl Biquad {
697    /// Create a new biquad with given coefficients.
698    pub fn new(b: [f64; 3], a: [f64; 2]) -> Self {
699        Self {
700            b,
701            a,
702            state: [0.0; 2],
703        }
704    }
705
706    /// Process a single sample (direct form II transposed).
707    pub fn process_sample(&mut self, x: f64) -> f64 {
708        let y = self.b[0] * x + self.state[0];
709        self.state[0] = self.b[1] * x - self.a[0] * y + self.state[1];
710        self.state[1] = self.b[2] * x - self.a[1] * y;
711        y
712    }
713
714    /// Process a block of samples.
715    pub fn process(&mut self, signal: &[f64]) -> Vec<f64> {
716        signal.iter().map(|&x| self.process_sample(x)).collect()
717    }
718
719    /// Reset filter state.
720    pub fn reset(&mut self) {
721        self.state = [0.0; 2];
722    }
723}
724
725/// A cascade of biquad sections forming an IIR filter.
726#[derive(Clone, Debug)]
727pub struct IirFilter {
728    /// Second-order sections.
729    pub sections: Vec<Biquad>,
730    /// Overall gain.
731    pub gain: f64,
732}
733
734impl IirFilter {
735    /// Create from a list of biquad sections and an overall gain.
736    pub fn new(sections: Vec<Biquad>, gain: f64) -> Self {
737        Self { sections, gain }
738    }
739
740    /// Process a signal through the IIR filter.
741    pub fn process(&mut self, signal: &[f64]) -> Vec<f64> {
742        let mut out: Vec<f64> = signal.iter().map(|&x| x * self.gain).collect();
743        for section in self.sections.iter_mut() {
744            out = section.process(&out);
745        }
746        out
747    }
748
749    /// Reset all section states.
750    pub fn reset(&mut self) {
751        for s in self.sections.iter_mut() {
752            s.reset();
753        }
754    }
755}
756
757// ---------------------------------------------------------------------------
758// Butterworth filter design
759// ---------------------------------------------------------------------------
760
761/// Butterworth low-pass filter design (bilinear transform, biquad cascade).
762///
763/// `order`: filter order. `wn`: normalised cutoff (0..1, 1 = Nyquist).
764pub fn butterworth_lowpass(order: usize, wn: f64) -> IirFilter {
765    // Pre-warp cutoff frequency
766    let wc = 2.0 * (PI * wn / 2.0).tan();
767    let n = order;
768    let mut sections = Vec::new();
769    let n_sections = n / 2;
770    // Pair up poles (and handle odd order separately)
771    for k in 0..n_sections {
772        let theta = PI * (2 * k + n + 1) as f64 / (2 * n) as f64;
773        let real = theta.cos();
774        let imag = theta.sin().abs();
775        // Analog prototype pole: s = wc * (real + j*imag)
776        // Bilinear transform: z = (1 + s/2) / (1 - s/2) -> use Tustin
777        let sr = wc * real;
778        let si = wc * imag;
779        // Bilinear: s -> 2*(z-1)/(z+1)
780        // Denominator: (z-1)^2 - (s/2)*(z+1)^2 ...
781        // Direct analog-to-digital via bilinear for second-order:
782        // H(s) = wc^2 / (s^2 - 2*sr*s + sr^2 + si^2)
783        let omega0_sq = sr * sr + si * si;
784        let two_alpha = -2.0 * sr;
785        // Bilinear coefficients for second-order LP section
786        let c = 2.0; // fs normalised
787        let b0d = omega0_sq;
788        let b1d = 2.0 * omega0_sq;
789        let b2d = omega0_sq;
790        let a0d = c * c + two_alpha * c + omega0_sq;
791        let a1d = -2.0 * c * c + 2.0 * omega0_sq;
792        let a2d = c * c - two_alpha * c + omega0_sq;
793        let bq = Biquad::new([b0d / a0d, b1d / a0d, b2d / a0d], [a1d / a0d, a2d / a0d]);
794        sections.push(bq);
795    }
796    // Handle odd order (first-order section)
797    let mut gain = 1.0;
798    if n % 2 == 1 {
799        // First-order LP: H(s) = wc / (s + wc)
800        // Bilinear: H(z) = (wc/2 + wc/2 z^{-1}) / ((1 + wc/2) + (-1 + wc/2) z^{-1})
801        let c = 1.0;
802        let b0 = wc / (wc + c);
803        let b1 = wc / (wc + c);
804        let a1 = (wc - c) / (wc + c);
805        let bq = Biquad::new([b0, b1, 0.0], [a1, 0.0]);
806        sections.push(bq);
807        gain = 1.0;
808    }
809    IirFilter::new(sections, gain)
810}
811
812/// Compute the magnitude response of a cascade of biquads at frequency `f` (normalised 0..1).
813pub fn iir_magnitude_response(filter: &IirFilter, f: f64) -> f64 {
814    let z = Complex64::exp_i(PI * f);
815    let mut h = Complex64::new(filter.gain, 0.0);
816    for bq in &filter.sections {
817        let num = Complex64::new(bq.b[0], 0.0)
818            + Complex64::new(bq.b[1], 0.0) * z.conj()
819            + Complex64::new(bq.b[2], 0.0) * z.conj() * z.conj();
820        let den = Complex64::new(1.0, 0.0)
821            + Complex64::new(bq.a[0], 0.0) * z.conj()
822            + Complex64::new(bq.a[1], 0.0) * z.conj() * z.conj();
823        // Avoid division by zero
824        if den.abs() > 1e-15 {
825            let r = den.norm_sq();
826            let den_inv = Complex64::new(den.re / r, -den.im / r);
827            h = h * num * den_inv;
828        }
829    }
830    h.abs()
831}
832
833// ---------------------------------------------------------------------------
834// Chebyshev Type I filter design
835// ---------------------------------------------------------------------------
836
837/// Chebyshev Type I low-pass filter (biquad cascade).
838///
839/// `order`: filter order. `wn`: normalised cutoff (0..1). `ripple_db`: passband ripple in dB.
840pub fn chebyshev1_lowpass(order: usize, wn: f64, ripple_db: f64) -> IirFilter {
841    let eps = (10.0f64.powf(ripple_db / 10.0) - 1.0).sqrt();
842    let wc = 2.0 * (PI * wn / 2.0).tan();
843    let n = order as f64;
844    let asinh_inv_eps = (1.0 / eps).asinh();
845    let mut sections = Vec::new();
846
847    let n_sections = order / 2;
848    for k in 0..n_sections {
849        let theta_k = PI * (2 * k + 1) as f64 / (2 * order) as f64;
850        let sigma = -(asinh_inv_eps / n).sinh() * theta_k.sin();
851        let omega = (asinh_inv_eps / n).cosh() * theta_k.cos();
852        let sr = wc * sigma;
853        let si = wc * omega.abs();
854        let omega0_sq = sr * sr + si * si;
855        let two_alpha = -2.0 * sr;
856        let c = 2.0;
857        let a0d = c * c + two_alpha * c + omega0_sq;
858        let a1d = -2.0 * c * c + 2.0 * omega0_sq;
859        let a2d = c * c - two_alpha * c + omega0_sq;
860        let bq = Biquad::new(
861            [omega0_sq / a0d, 2.0 * omega0_sq / a0d, omega0_sq / a0d],
862            [a1d / a0d, a2d / a0d],
863        );
864        sections.push(bq);
865    }
866    if order % 2 == 1 {
867        let sigma0 = -(asinh_inv_eps / n).sinh();
868        let sr = wc * sigma0;
869        let c = 1.0;
870        let b0 = (-sr) / (c - sr);
871        let b1 = (-sr) / (c - sr);
872        let a1 = (-sr - c) / (c - sr);
873        let bq = Biquad::new([b0, b1, 0.0], [a1, 0.0]);
874        sections.push(bq);
875    }
876    IirFilter::new(sections, 1.0)
877}
878
879// ---------------------------------------------------------------------------
880// Hilbert Transform
881// ---------------------------------------------------------------------------
882
883/// Compute the analytic signal (x + j*H{x}) using FFT.
884///
885/// Returns the complex analytic signal from which the instantaneous
886/// amplitude and frequency can be derived.
887pub fn analytic_signal(x: &[f64]) -> Vec<Complex64> {
888    let n = x.len().next_power_of_two();
889    let mut spec = fft(x);
890    spec.resize(n, Complex64::new(0.0, 0.0));
891    // Zero the negative frequencies and double positive ones
892    for k in 1..n / 2 {
893        spec[k] = spec[k].scale(2.0);
894    }
895    for k in n / 2 + 1..n {
896        spec[k] = Complex64::new(0.0, 0.0);
897    }
898    let mut result = spec;
899    ifft_inplace(&mut result);
900    result[..x.len()].to_vec()
901}
902
903/// Compute the Hilbert transform of a real signal.
904pub fn hilbert_transform(x: &[f64]) -> Vec<f64> {
905    analytic_signal(x).iter().map(|c| c.im).collect()
906}
907
908/// Compute the instantaneous amplitude (envelope) of a signal.
909pub fn instantaneous_amplitude(x: &[f64]) -> Vec<f64> {
910    analytic_signal(x).iter().map(|c| c.abs()).collect()
911}
912
913/// Compute the instantaneous phase of a signal (radians).
914pub fn instantaneous_phase(x: &[f64]) -> Vec<f64> {
915    analytic_signal(x).iter().map(|c| c.arg()).collect()
916}
917
918/// Compute the instantaneous frequency (Hz) given sample rate `fs`.
919pub fn instantaneous_frequency(x: &[f64], fs: f64) -> Vec<f64> {
920    let phase = instantaneous_phase(x);
921    let mut freq = vec![0.0f64; phase.len()];
922    for i in 1..phase.len() {
923        let mut dphi = phase[i] - phase[i - 1];
924        // Unwrap
925        while dphi > PI {
926            dphi -= 2.0 * PI;
927        }
928        while dphi < -PI {
929            dphi += 2.0 * PI;
930        }
931        freq[i] = dphi * fs / (2.0 * PI);
932    }
933    freq[0] = freq[1];
934    freq
935}
936
937// ---------------------------------------------------------------------------
938// Spectrogram
939// ---------------------------------------------------------------------------
940
941/// Spectrogram result: time bins, frequency bins, and STFT magnitude matrix.
942#[derive(Clone, Debug)]
943pub struct Spectrogram {
944    /// Time bin centres (seconds).
945    pub times: Vec<f64>,
946    /// Frequency bins (Hz).
947    pub freqs: Vec<f64>,
948    /// Magnitude\[time\]\[freq\].
949    pub magnitude: Vec<Vec<f64>>,
950    /// Phase\[time\]\[freq\] (radians).
951    pub phase: Vec<Vec<f64>>,
952}
953
954/// Compute a short-time Fourier transform (STFT) spectrogram.
955pub fn spectrogram(
956    signal: &[f64],
957    fs: f64,
958    window_len: usize,
959    hop: usize,
960    window_type: WindowType,
961) -> Spectrogram {
962    let window = match window_type {
963        WindowType::Hann => hann_window(window_len),
964        WindowType::Hamming => hamming_window(window_len),
965        WindowType::Blackman => blackman_window(window_len),
966        _ => hann_window(window_len),
967    };
968    let n_fft = window_len.next_power_of_two();
969    let n_bins = n_fft / 2 + 1;
970
971    let mut times = Vec::new();
972    let mut magnitudes = Vec::new();
973    let mut phases = Vec::new();
974
975    let mut start = 0;
976    while start + window_len <= signal.len() {
977        times.push((start + window_len / 2) as f64 / fs);
978        let mut buf: Vec<f64> = signal[start..start + window_len]
979            .iter()
980            .zip(window.iter())
981            .map(|(s, w)| s * w)
982            .collect();
983        buf.resize(n_fft, 0.0);
984        let spec = fft(&buf);
985        magnitudes.push(spec[..n_bins].iter().map(|c| c.abs()).collect::<Vec<_>>());
986        phases.push(spec[..n_bins].iter().map(|c| c.arg()).collect::<Vec<_>>());
987        start += hop;
988    }
989
990    let freqs = rfft_frequencies(n_fft, fs);
991    Spectrogram {
992        times,
993        freqs,
994        magnitude: magnitudes,
995        phase: phases,
996    }
997}
998
999// ---------------------------------------------------------------------------
1000// Wavelet Transforms
1001// ---------------------------------------------------------------------------
1002
1003/// Haar wavelet decomposition (one level).
1004///
1005/// Returns (approximation, detail) coefficient vectors.
1006pub fn haar_decompose(signal: &[f64]) -> (Vec<f64>, Vec<f64>) {
1007    let n = signal.len() / 2;
1008    let sqrt2_inv = 1.0 / 2.0f64.sqrt();
1009    let approx: Vec<f64> = (0..n)
1010        .map(|i| (signal[2 * i] + signal[2 * i + 1]) * sqrt2_inv)
1011        .collect();
1012    let detail: Vec<f64> = (0..n)
1013        .map(|i| (signal[2 * i] - signal[2 * i + 1]) * sqrt2_inv)
1014        .collect();
1015    (approx, detail)
1016}
1017
1018/// Haar wavelet reconstruction (one level).
1019pub fn haar_reconstruct(approx: &[f64], detail: &[f64]) -> Vec<f64> {
1020    let n = approx.len();
1021    let sqrt2_inv = 1.0 / 2.0f64.sqrt();
1022    let mut out = vec![0.0f64; 2 * n];
1023    for i in 0..n {
1024        out[2 * i] = (approx[i] + detail[i]) * sqrt2_inv;
1025        out[2 * i + 1] = (approx[i] - detail[i]) * sqrt2_inv;
1026    }
1027    out
1028}
1029
1030/// Multi-level Haar DWT.
1031///
1032/// Returns a flattened coefficient array: \[level_J_approx | details_J | details_{J-1} | ... | details_1\].
1033pub fn haar_dwt(signal: &[f64], levels: usize) -> Vec<f64> {
1034    let mut approx = signal.to_vec();
1035    let mut details_stack: Vec<Vec<f64>> = Vec::new();
1036    for _ in 0..levels {
1037        if approx.len() < 2 {
1038            break;
1039        }
1040        let (a, d) = haar_decompose(&approx);
1041        approx = a;
1042        details_stack.push(d);
1043    }
1044    let mut result = approx;
1045    for d in details_stack.into_iter().rev() {
1046        result.extend(d);
1047    }
1048    result
1049}
1050
1051/// Multi-level Haar inverse DWT.
1052pub fn haar_idwt(coeffs: &[f64], levels: usize) -> Vec<f64> {
1053    let total = coeffs.len();
1054    let approx_len = total >> levels;
1055    let mut approx = coeffs[..approx_len].to_vec();
1056    let mut pos = approx_len;
1057    for _ in 0..levels {
1058        let detail_len = approx.len();
1059        if pos + detail_len > coeffs.len() {
1060            break;
1061        }
1062        let detail = &coeffs[pos..pos + detail_len];
1063        approx = haar_reconstruct(&approx, detail);
1064        pos += detail_len;
1065    }
1066    approx
1067}
1068
1069// Daubechies db2 (Daubechies-4) filter coefficients
1070const DB2_LO: [f64; 4] = [0.6830127, 1.1830127, 0.3169873, -0.1830127];
1071const DB2_HI: [f64; 4] = [-0.1830127, -0.3169873, 1.1830127, -0.6830127];
1072// Inverse (synthesis) low/high pass
1073const DB2_LO_R: [f64; 4] = [-0.1830127, 0.3169873, 1.1830127, 0.6830127];
1074const DB2_HI_R: [f64; 4] = [-0.6830127, 1.1830127, -0.3169873, -0.1830127];
1075
1076/// Daubechies-4 (db2) single-level forward DWT.
1077pub fn db2_decompose(signal: &[f64]) -> (Vec<f64>, Vec<f64>) {
1078    let n = signal.len();
1079    let filt_len = DB2_LO.len();
1080    let out_len = (n + filt_len - 1) / 2;
1081    let mut approx = vec![0.0f64; out_len];
1082    let mut detail = vec![0.0f64; out_len];
1083    for k in 0..out_len {
1084        let mut a = 0.0;
1085        let mut d = 0.0;
1086        for j in 0..filt_len {
1087            let idx = 2 * k + j;
1088            let s = if idx < n { signal[idx] } else { 0.0 };
1089            a += DB2_LO[j] * s;
1090            d += DB2_HI[j] * s;
1091        }
1092        approx[k] = a / 2.0f64.sqrt();
1093        detail[k] = d / 2.0f64.sqrt();
1094    }
1095    (approx, detail)
1096}
1097
1098/// Daubechies-4 (db2) single-level inverse DWT.
1099pub fn db2_reconstruct(approx: &[f64], detail: &[f64]) -> Vec<f64> {
1100    let n = approx.len();
1101    let out_len = 2 * n;
1102    let mut out = vec![0.0f64; out_len + DB2_LO_R.len()];
1103    for k in 0..n {
1104        for j in 0..DB2_LO_R.len() {
1105            let idx = 2 * k + j;
1106            if idx < out.len() {
1107                out[idx] += approx[k] * DB2_LO_R[j] / 2.0f64.sqrt();
1108                out[idx] += detail[k] * DB2_HI_R[j] / 2.0f64.sqrt();
1109            }
1110        }
1111    }
1112    out[..out_len].to_vec()
1113}
1114
1115// ---------------------------------------------------------------------------
1116// Resampling
1117// ---------------------------------------------------------------------------
1118
1119/// Upsample a signal by integer factor `up` (insert zeros between samples).
1120pub fn upsample(signal: &[f64], up: usize) -> Vec<f64> {
1121    let mut out = vec![0.0f64; signal.len() * up];
1122    for (i, &s) in signal.iter().enumerate() {
1123        out[i * up] = s;
1124    }
1125    out
1126}
1127
1128/// Downsample a signal by integer factor `down` (decimate, no anti-aliasing).
1129pub fn downsample(signal: &[f64], down: usize) -> Vec<f64> {
1130    signal.iter().step_by(down).cloned().collect()
1131}
1132
1133/// Resample a signal from sample rate `fs_in` to `fs_out` using rational resampling.
1134///
1135/// Computes the integer up/down ratio and applies a low-pass anti-aliasing filter.
1136pub fn resample(signal: &[f64], fs_in: f64, fs_out: f64) -> Vec<f64> {
1137    // Compute rational ratio using GCD approximation
1138    let ratio = fs_out / fs_in;
1139    let up = (ratio * 100.0).round() as usize;
1140    let down = 100usize;
1141    let up_factor = up.max(1);
1142    let down_factor = down.max(1);
1143    // Upsample
1144    let up_sig = upsample(signal, up_factor);
1145    // Anti-aliasing filter
1146    let cutoff = 1.0 / up_factor.max(down_factor) as f64;
1147    let filt = FirFilter::lowpass(65, cutoff);
1148    let filtered = filt.apply(&up_sig);
1149    // Downsample
1150
1151    downsample(&filtered, down_factor)
1152}
1153
1154/// Linear interpolation resampling.
1155pub fn resample_linear(signal: &[f64], n_out: usize) -> Vec<f64> {
1156    if signal.is_empty() || n_out == 0 {
1157        return Vec::new();
1158    }
1159    let n_in = signal.len();
1160    (0..n_out)
1161        .map(|i| {
1162            let t = i as f64 * (n_in - 1) as f64 / (n_out - 1).max(1) as f64;
1163            let lo = t as usize;
1164            let hi = (lo + 1).min(n_in - 1);
1165            let frac = t - lo as f64;
1166            signal[lo] * (1.0 - frac) + signal[hi] * frac
1167        })
1168        .collect()
1169}
1170
1171// ---------------------------------------------------------------------------
1172// Utility helpers
1173// ---------------------------------------------------------------------------
1174
1175/// Compute the RMS value of a signal.
1176pub fn rms(signal: &[f64]) -> f64 {
1177    let n = signal.len() as f64;
1178    (signal.iter().map(|x| x * x).sum::<f64>() / n).sqrt()
1179}
1180
1181/// Compute the mean of a signal.
1182pub fn signal_mean(signal: &[f64]) -> f64 {
1183    signal.iter().sum::<f64>() / signal.len() as f64
1184}
1185
1186/// Compute the variance of a signal.
1187pub fn signal_variance(signal: &[f64]) -> f64 {
1188    let mean = signal_mean(signal);
1189    signal.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / signal.len() as f64
1190}
1191
1192/// Compute the SNR (signal-to-noise ratio) in dB given signal and noise power.
1193pub fn snr_db(signal_power: f64, noise_power: f64) -> f64 {
1194    10.0 * (signal_power / noise_power.max(1e-300)).log10()
1195}
1196
1197/// Compute the THD (total harmonic distortion) from a power spectrum.
1198///
1199/// `fundamental_bin`: FFT bin index of the fundamental frequency.
1200/// `n_harmonics`: number of harmonics to include.
1201pub fn total_harmonic_distortion(
1202    power_spec: &[f64],
1203    fundamental_bin: usize,
1204    n_harmonics: usize,
1205) -> f64 {
1206    let p1 = power_spec[fundamental_bin];
1207    if p1 == 0.0 {
1208        return 0.0;
1209    }
1210    let harmonic_power: f64 = (2..=n_harmonics)
1211        .filter_map(|h| {
1212            let bin = fundamental_bin * h;
1213            if bin < power_spec.len() {
1214                Some(power_spec[bin])
1215            } else {
1216                None
1217            }
1218        })
1219        .sum();
1220    (harmonic_power / p1).sqrt()
1221}
1222
1223/// Zero-pad a signal to length `n`.
1224pub fn zero_pad(signal: &[f64], n: usize) -> Vec<f64> {
1225    let mut out = signal.to_vec();
1226    out.resize(n, 0.0);
1227    out
1228}
1229
1230/// Unwrap phase (correct for 2π discontinuities).
1231pub fn unwrap_phase(phase: &[f64]) -> Vec<f64> {
1232    let mut out = phase.to_vec();
1233    for i in 1..out.len() {
1234        let diff = out[i] - out[i - 1];
1235        if diff > PI {
1236            for v in out[i..].iter_mut() {
1237                *v -= 2.0 * PI;
1238            }
1239        } else if diff < -PI {
1240            for v in out[i..].iter_mut() {
1241                *v += 2.0 * PI;
1242            }
1243        }
1244    }
1245    out
1246}
1247
1248/// Compute group delay from the phase response.
1249pub fn group_delay(phase: &[f64], df: f64) -> Vec<f64> {
1250    let n = phase.len();
1251    let mut gd = vec![0.0f64; n];
1252    for i in 1..n - 1 {
1253        gd[i] = -(phase[i + 1] - phase[i - 1]) / (2.0 * df);
1254    }
1255    gd[0] = gd[1];
1256    gd[n - 1] = gd[n - 2];
1257    gd
1258}
1259
1260/// Compute the energy of a signal.
1261pub fn signal_energy(signal: &[f64]) -> f64 {
1262    signal.iter().map(|x| x * x).sum()
1263}
1264
1265/// Normalize a signal to unit energy.
1266pub fn normalize_energy(signal: &[f64]) -> Vec<f64> {
1267    let e = signal_energy(signal).sqrt().max(1e-300);
1268    signal.iter().map(|x| x / e).collect()
1269}
1270
1271/// Normalize a signal to unit peak amplitude.
1272pub fn normalize_peak(signal: &[f64]) -> Vec<f64> {
1273    let peak = signal
1274        .iter()
1275        .map(|x| x.abs())
1276        .fold(0.0f64, f64::max)
1277        .max(1e-300);
1278    signal.iter().map(|x| x / peak).collect()
1279}
1280
1281/// Generate a test sinusoidal signal.
1282pub fn generate_sine(freq: f64, fs: f64, n_samples: usize, amplitude: f64, phase: f64) -> Vec<f64> {
1283    (0..n_samples)
1284        .map(|i| amplitude * (2.0 * PI * freq * i as f64 / fs + phase).sin())
1285        .collect()
1286}
1287
1288/// Generate a test chirp (frequency sweep) signal.
1289pub fn generate_chirp(f0: f64, f1: f64, fs: f64, n_samples: usize, amplitude: f64) -> Vec<f64> {
1290    let duration = n_samples as f64 / fs;
1291    (0..n_samples)
1292        .map(|i| {
1293            let t = i as f64 / fs;
1294            let inst_freq = f0 + (f1 - f0) * t / duration;
1295            amplitude * (2.0 * PI * inst_freq * t).sin()
1296        })
1297        .collect()
1298}
1299
1300/// Generate white Gaussian noise (Box-Muller transform).
1301pub fn generate_white_noise(n: usize, amplitude: f64, seed: u64) -> Vec<f64> {
1302    let mut x = seed as f64 + 1.0;
1303    let mut noise = Vec::with_capacity(n);
1304    for _ in 0..n / 2 {
1305        // LCG-based pseudo-random
1306        x = (x * 6364136223846793005.0 + 1442695040888963407.0).abs() % 1e18;
1307        let u1 = x / 1e18;
1308        x = (x * 6364136223846793005.0 + 1442695040888963407.0).abs() % 1e18;
1309        let u2 = x / 1e18;
1310        let r = (-2.0 * (u1 + 1e-300).ln()).sqrt();
1311        noise.push(amplitude * r * (2.0 * PI * u2).cos());
1312        noise.push(amplitude * r * (2.0 * PI * u2).sin());
1313    }
1314    if n % 2 == 1 {
1315        noise.push(0.0);
1316    }
1317    noise
1318}
1319
1320// ---------------------------------------------------------------------------
1321// Peak detection
1322// ---------------------------------------------------------------------------
1323
1324/// Detect local maxima in a signal.
1325///
1326/// Returns the indices of peaks where the value is greater than both neighbours.
1327pub fn find_peaks(signal: &[f64]) -> Vec<usize> {
1328    (1..signal.len() - 1)
1329        .filter(|&i| signal[i] > signal[i - 1] && signal[i] > signal[i + 1])
1330        .collect()
1331}
1332
1333/// Detect peaks above a threshold.
1334pub fn find_peaks_above(signal: &[f64], threshold: f64) -> Vec<usize> {
1335    find_peaks(signal)
1336        .into_iter()
1337        .filter(|&i| signal[i] > threshold)
1338        .collect()
1339}
1340
1341// ---------------------------------------------------------------------------
1342// Moving statistics
1343// ---------------------------------------------------------------------------
1344
1345/// Compute the moving average of a signal.
1346pub fn moving_average(signal: &[f64], window: usize) -> Vec<f64> {
1347    let mut out = Vec::with_capacity(signal.len());
1348    let mut sum = 0.0;
1349    for (i, &x) in signal.iter().enumerate() {
1350        sum += x;
1351        if i >= window {
1352            sum -= signal[i - window];
1353        }
1354        out.push(sum / window.min(i + 1) as f64);
1355    }
1356    out
1357}
1358
1359/// Compute the moving RMS of a signal.
1360pub fn moving_rms(signal: &[f64], window: usize) -> Vec<f64> {
1361    let mut out = Vec::with_capacity(signal.len());
1362    let mut sum_sq = 0.0;
1363    for (i, &x) in signal.iter().enumerate() {
1364        sum_sq += x * x;
1365        if i >= window {
1366            sum_sq -= signal[i - window].powi(2);
1367        }
1368        out.push((sum_sq / window.min(i + 1) as f64).sqrt());
1369    }
1370    out
1371}
1372
1373// ---------------------------------------------------------------------------
1374// Cepstrum
1375// ---------------------------------------------------------------------------
1376
1377/// Compute the real cepstrum of a signal.
1378pub fn real_cepstrum(signal: &[f64]) -> Vec<f64> {
1379    let spec = fft(signal);
1380    let log_mag: Vec<f64> = spec.iter().map(|c| c.abs().max(1e-300).ln()).collect();
1381    // IFFT of log magnitude (treat as symmetric real spectrum)
1382    let mut buf: Vec<Complex64> = log_mag.iter().map(|&v| Complex64::new(v, 0.0)).collect();
1383    ifft_inplace(&mut buf);
1384    buf.iter().map(|c| c.re).collect()
1385}
1386
1387/// Compute the complex cepstrum of a signal.
1388pub fn complex_cepstrum(signal: &[f64]) -> Vec<f64> {
1389    let spec = fft(signal);
1390    let log_spec: Vec<Complex64> = spec
1391        .iter()
1392        .map(|c| {
1393            let r = c.abs().max(1e-300);
1394            Complex64::new(r.ln(), c.arg())
1395        })
1396        .collect();
1397    let mut buf = log_spec;
1398    ifft_inplace(&mut buf);
1399    buf.iter().map(|c| c.re).collect()
1400}
1401
1402// ---------------------------------------------------------------------------
1403// STFT and inverse STFT
1404// ---------------------------------------------------------------------------
1405
1406/// Short-time Fourier transform: returns complex coefficients \[frame\]\[freq_bin\].
1407pub fn stft(
1408    signal: &[f64],
1409    window_len: usize,
1410    hop: usize,
1411    window_type: WindowType,
1412) -> Vec<Vec<Complex64>> {
1413    let window = match window_type {
1414        WindowType::Hann => hann_window(window_len),
1415        WindowType::Hamming => hamming_window(window_len),
1416        WindowType::Blackman => blackman_window(window_len),
1417        _ => hann_window(window_len),
1418    };
1419    let n_fft = window_len.next_power_of_two();
1420    let n_bins = n_fft / 2 + 1;
1421    let mut frames = Vec::new();
1422    let mut start = 0;
1423    while start + window_len <= signal.len() {
1424        let mut buf: Vec<f64> = signal[start..start + window_len]
1425            .iter()
1426            .zip(window.iter())
1427            .map(|(s, w)| s * w)
1428            .collect();
1429        buf.resize(n_fft, 0.0);
1430        let spec = fft(&buf);
1431        frames.push(spec[..n_bins].to_vec());
1432        start += hop;
1433    }
1434    frames
1435}
1436
1437/// Inverse STFT using overlap-add.
1438pub fn istft(
1439    frames: &[Vec<Complex64>],
1440    window_len: usize,
1441    hop: usize,
1442    window_type: WindowType,
1443) -> Vec<f64> {
1444    let window = match window_type {
1445        WindowType::Hann => hann_window(window_len),
1446        WindowType::Hamming => hamming_window(window_len),
1447        _ => hann_window(window_len),
1448    };
1449    let n_fft = window_len.next_power_of_two();
1450    let signal_len = hop * frames.len() + window_len;
1451    let mut out = vec![0.0f64; signal_len];
1452    let mut norm = vec![0.0f64; signal_len];
1453
1454    for (frame_idx, frame) in frames.iter().enumerate() {
1455        let mut spec = frame.to_vec();
1456        // Mirror for real IFFT
1457        while spec.len() < n_fft {
1458            let k = n_fft - spec.len();
1459            spec.push(frames[frame_idx][k.min(frame.len() - 1)].conj());
1460        }
1461        ifft_inplace(&mut spec);
1462        let start = frame_idx * hop;
1463        for (i, (&s, &w)) in spec.iter().zip(window.iter()).enumerate().take(window_len) {
1464            out[start + i] += s.re * w;
1465            norm[start + i] += w * w;
1466        }
1467    }
1468    for (o, n) in out.iter_mut().zip(norm.iter()) {
1469        if *n > 1e-10 {
1470            *o /= n;
1471        }
1472    }
1473    out
1474}
1475
1476// ---------------------------------------------------------------------------
1477// Tests
1478// ---------------------------------------------------------------------------
1479
1480#[cfg(test)]
1481mod tests {
1482    use super::*;
1483
1484    const EPS: f64 = 1e-9;
1485
1486    fn approx_eq(a: f64, b: f64, tol: f64) -> bool {
1487        (a - b).abs() < tol
1488    }
1489
1490    // --- Complex64 ---
1491
1492    #[test]
1493    fn test_complex_basic_ops() {
1494        let a = Complex64::new(3.0, 4.0);
1495        let b = Complex64::new(1.0, -2.0);
1496        let sum = a + b;
1497        assert_eq!(sum.re, 4.0);
1498        assert_eq!(sum.im, 2.0);
1499        let prod = a * b;
1500        // (3+4i)(1-2i) = 3 - 6i + 4i + 8 = 11 - 2i
1501        assert!(approx_eq(prod.re, 11.0, EPS));
1502        assert!(approx_eq(prod.im, -2.0, EPS));
1503    }
1504
1505    #[test]
1506    fn test_complex_abs_and_arg() {
1507        let c = Complex64::new(3.0, 4.0);
1508        assert!(approx_eq(c.abs(), 5.0, EPS));
1509    }
1510
1511    #[test]
1512    fn test_complex_conj() {
1513        let c = Complex64::new(1.0, -2.0);
1514        let cc = c.conj();
1515        assert_eq!(cc.re, 1.0);
1516        assert_eq!(cc.im, 2.0);
1517    }
1518
1519    #[test]
1520    fn test_complex_exp_i() {
1521        let c = Complex64::exp_i(0.0);
1522        assert!(approx_eq(c.re, 1.0, EPS));
1523        assert!(approx_eq(c.im, 0.0, EPS));
1524        let c90 = Complex64::exp_i(PI / 2.0);
1525        assert!(approx_eq(c90.re, 0.0, 1e-14));
1526        assert!(approx_eq(c90.im, 1.0, 1e-14));
1527    }
1528
1529    // --- FFT ---
1530
1531    #[test]
1532    fn test_fft_dc() {
1533        // DC signal: all ones -> FFT has only X[0] = N, rest 0
1534        let signal = vec![1.0f64; 8];
1535        let spec = fft(&signal);
1536        assert!(approx_eq(spec[0].re, 8.0, 1e-10));
1537        assert!(approx_eq(spec[0].im, 0.0, 1e-10));
1538        for k in 1..8 {
1539            assert!(approx_eq(spec[k].re, 0.0, 1e-10));
1540            assert!(approx_eq(spec[k].im, 0.0, 1e-10));
1541        }
1542    }
1543
1544    #[test]
1545    fn test_fft_single_tone() {
1546        // cos(2π*k*n/N) -> peaks at bin k and N-k
1547        let n = 16usize;
1548        let k = 2usize;
1549        let signal: Vec<f64> = (0..n)
1550            .map(|i| (2.0 * PI * k as f64 * i as f64 / n as f64).cos())
1551            .collect();
1552        let spec = fft(&signal);
1553        // Bin k should have magnitude ≈ N/2
1554        assert!(approx_eq(spec[k].abs(), n as f64 / 2.0, 1e-8));
1555    }
1556
1557    #[test]
1558    fn test_fft_ifft_roundtrip() {
1559        let signal = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
1560        let spec = fft(&signal);
1561        let recovered = ifft_real(&spec);
1562        for (a, b) in signal.iter().zip(recovered.iter()) {
1563            assert!(approx_eq(*a, *b, 1e-10));
1564        }
1565    }
1566
1567    #[test]
1568    fn test_fft_parseval() {
1569        // Parseval's theorem: sum |x[n]|^2 = (1/N) * sum |X[k]|^2
1570        let signal = vec![1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 0.5, -0.5];
1571        let n = signal.len() as f64;
1572        let time_energy: f64 = signal.iter().map(|x| x * x).sum();
1573        let spec = fft(&signal);
1574        let freq_energy: f64 = spec.iter().map(|c| c.norm_sq()).sum::<f64>() / n;
1575        assert!(approx_eq(time_energy, freq_energy, 1e-8));
1576    }
1577
1578    // --- Windows ---
1579
1580    #[test]
1581    fn test_hann_window_endpoints() {
1582        let w = hann_window(8);
1583        assert!(approx_eq(w[0], 0.0, 1e-10));
1584    }
1585
1586    #[test]
1587    fn test_hamming_window_range() {
1588        let w = hamming_window(64);
1589        for v in &w {
1590            assert!(*v >= 0.0 && *v <= 1.0);
1591        }
1592    }
1593
1594    #[test]
1595    fn test_blackman_window_endpoints() {
1596        let w = blackman_window(16);
1597        assert!(w[0].abs() < 1e-10);
1598    }
1599
1600    #[test]
1601    fn test_kaiser_window_unity_center() {
1602        let w = kaiser_window(9, 4.0);
1603        // Center should be 1.0
1604        assert!(approx_eq(w[4], 1.0, 1e-10));
1605    }
1606
1607    #[test]
1608    fn test_bessel_i0() {
1609        // I0(0) = 1
1610        assert!(approx_eq(bessel_i0(0.0), 1.0, EPS));
1611        // I0(1) ≈ 1.26606
1612        assert!(approx_eq(bessel_i0(1.0), 1.2660658, 1e-6));
1613    }
1614
1615    // --- Convolution ---
1616
1617    #[test]
1618    fn test_convolve_identity() {
1619        let x = vec![1.0, 2.0, 3.0];
1620        let h = vec![1.0, 0.0, 0.0];
1621        let y = convolve(&x, &h);
1622        assert!(approx_eq(y[0], 1.0, 1e-10));
1623        assert!(approx_eq(y[1], 2.0, 1e-10));
1624        assert!(approx_eq(y[2], 3.0, 1e-10));
1625    }
1626
1627    #[test]
1628    fn test_convolve_box() {
1629        // Convolve [1,1,1] with [1,1] -> [1,2,2,1]
1630        let x = vec![1.0, 1.0, 1.0];
1631        let h = vec![1.0, 1.0];
1632        let y = convolve(&x, &h);
1633        assert!(approx_eq(y[0], 1.0, 1e-10));
1634        assert!(approx_eq(y[1], 2.0, 1e-10));
1635        assert!(approx_eq(y[2], 2.0, 1e-10));
1636        assert!(approx_eq(y[3], 1.0, 1e-10));
1637    }
1638
1639    #[test]
1640    fn test_auto_correlate_peak() {
1641        let signal = vec![1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0, 0.0];
1642        let ac = auto_correlate(&signal);
1643        // Peak should be at lag = signal.len() - 1
1644        let mid = signal.len() - 1;
1645        let peak_idx = ac
1646            .iter()
1647            .enumerate()
1648            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
1649            .map(|(i, _)| i)
1650            .unwrap();
1651        assert_eq!(peak_idx, mid);
1652    }
1653
1654    // --- FIR filter ---
1655
1656    #[test]
1657    fn test_fir_lowpass_dc_gain() {
1658        let fir = FirFilter::lowpass(31, 0.5);
1659        // DC gain should be ≈ 1.0
1660        let dc_gain: f64 = fir.coeffs.iter().sum();
1661        assert!(approx_eq(dc_gain, 1.0, 0.05));
1662    }
1663
1664    #[test]
1665    fn test_fir_highpass_dc_reject() {
1666        let fir = FirFilter::highpass(31, 0.5);
1667        let dc_gain: f64 = fir.coeffs.iter().sum();
1668        assert!(dc_gain.abs() < 0.1);
1669    }
1670
1671    #[test]
1672    fn test_fir_apply_length() {
1673        let fir = FirFilter::lowpass(15, 0.3);
1674        let x = vec![1.0; 64];
1675        let y = fir.apply(&x);
1676        assert_eq!(y.len(), 64 + 15 - 1);
1677    }
1678
1679    // --- Biquad IIR ---
1680
1681    #[test]
1682    fn test_biquad_passthrough() {
1683        // Unity gain: b=[1,0,0], a=[0,0]
1684        let mut bq = Biquad::new([1.0, 0.0, 0.0], [0.0, 0.0]);
1685        let x = vec![1.0, 2.0, 3.0, 4.0];
1686        let y = bq.process(&x);
1687        assert!(approx_eq(y[0], 1.0, EPS));
1688        assert!(approx_eq(y[1], 2.0, EPS));
1689    }
1690
1691    #[test]
1692    fn test_biquad_reset() {
1693        let mut bq = Biquad::new([0.5, 0.5, 0.0], [0.0, 0.0]);
1694        let _ = bq.process(&[1.0, 1.0, 1.0]);
1695        bq.reset();
1696        let y = bq.process(&[1.0]);
1697        assert!(approx_eq(y[0], 0.5, EPS));
1698    }
1699
1700    // --- Hilbert transform ---
1701
1702    #[test]
1703    fn test_hilbert_of_cos_is_sin() {
1704        let n = 256;
1705        let fs = 256.0f64;
1706        let f = 10.0f64;
1707        let cos_sig: Vec<f64> = (0..n)
1708            .map(|i| (2.0 * PI * f * i as f64 / fs).cos())
1709            .collect();
1710        let sin_sig: Vec<f64> = (0..n)
1711            .map(|i| (2.0 * PI * f * i as f64 / fs).sin())
1712            .collect();
1713        let ht = hilbert_transform(&cos_sig);
1714        // Compare middle section to avoid edge effects
1715        for i in 16..n - 16 {
1716            assert!(approx_eq(ht[i], sin_sig[i], 0.02));
1717        }
1718    }
1719
1720    #[test]
1721    fn test_instantaneous_amplitude() {
1722        let n = 256;
1723        let fs = 256.0f64;
1724        let f = 10.0f64;
1725        let amp = 3.0f64;
1726        let sig: Vec<f64> = (0..n)
1727            .map(|i| amp * (2.0 * PI * f * i as f64 / fs).cos())
1728            .collect();
1729        let env = instantaneous_amplitude(&sig);
1730        // Middle section should be ≈ amp
1731        for i in 16..n - 16 {
1732            assert!(approx_eq(env[i], amp, 0.1));
1733        }
1734    }
1735
1736    // --- Haar DWT ---
1737
1738    #[test]
1739    fn test_haar_decompose_reconstruct() {
1740        let signal = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
1741        let (approx, detail) = haar_decompose(&signal);
1742        let recon = haar_reconstruct(&approx, &detail);
1743        for (a, b) in signal.iter().zip(recon.iter()) {
1744            assert!(approx_eq(*a, *b, 1e-10));
1745        }
1746    }
1747
1748    #[test]
1749    fn test_haar_dwt_energy_preservation() {
1750        let signal: Vec<f64> = (0..8).map(|i| i as f64 + 1.0).collect();
1751        let coeffs = haar_dwt(&signal, 3);
1752        let orig_energy: f64 = signal.iter().map(|x| x * x).sum();
1753        let coeff_energy: f64 = coeffs.iter().map(|x| x * x).sum();
1754        assert!(approx_eq(orig_energy, coeff_energy, 1e-8));
1755    }
1756
1757    // --- Spectrogram ---
1758
1759    #[test]
1760    fn test_spectrogram_shape() {
1761        let n = 256;
1762        let fs = 1000.0f64;
1763        let signal: Vec<f64> = (0..n)
1764            .map(|i| (2.0 * PI * 100.0 * i as f64 / fs).sin())
1765            .collect();
1766        let sg = spectrogram(&signal, fs, 32, 16, WindowType::Hann);
1767        assert!(!sg.times.is_empty());
1768        assert!(!sg.freqs.is_empty());
1769        for frame in &sg.magnitude {
1770            assert_eq!(frame.len(), sg.freqs.len());
1771        }
1772    }
1773
1774    // --- PSD ---
1775
1776    #[test]
1777    fn test_welch_psd_bins_count() {
1778        let signal: Vec<f64> = (0..512).map(|i| (i as f64 * 0.1).sin()).collect();
1779        let result = welch_psd(&signal, 1000.0, 128, 64, WindowType::Hann);
1780        assert!(!result.psd.is_empty());
1781        assert_eq!(result.freqs.len(), result.psd.len());
1782    }
1783
1784    #[test]
1785    fn test_periodogram_length() {
1786        let signal = vec![0.0f64; 64];
1787        let psd = periodogram_psd(&signal, 100.0, WindowType::Hamming);
1788        assert!(!psd.psd.is_empty());
1789    }
1790
1791    // --- Resampling ---
1792
1793    #[test]
1794    fn test_resample_linear_length() {
1795        let signal = vec![0.0, 1.0, 2.0, 3.0];
1796        let out = resample_linear(&signal, 8);
1797        assert_eq!(out.len(), 8);
1798    }
1799
1800    #[test]
1801    fn test_upsample_downsample() {
1802        let signal = vec![1.0, 2.0, 3.0, 4.0];
1803        let up = upsample(&signal, 2);
1804        assert_eq!(up.len(), 8);
1805        assert_eq!(up[0], 1.0);
1806        assert_eq!(up[2], 2.0);
1807        let dn = downsample(&up, 2);
1808        assert_eq!(dn, signal);
1809    }
1810
1811    // --- Utilities ---
1812
1813    #[test]
1814    fn test_rms() {
1815        let signal = vec![1.0, -1.0, 1.0, -1.0];
1816        assert!(approx_eq(rms(&signal), 1.0, EPS));
1817    }
1818
1819    #[test]
1820    fn test_signal_energy() {
1821        let signal = vec![3.0, 4.0];
1822        assert!(approx_eq(signal_energy(&signal), 25.0, EPS));
1823    }
1824
1825    #[test]
1826    fn test_normalize_energy() {
1827        let signal = vec![3.0, 4.0];
1828        let normed = normalize_energy(&signal);
1829        assert!(approx_eq(signal_energy(&normed), 1.0, 1e-10));
1830    }
1831
1832    #[test]
1833    fn test_find_peaks() {
1834        let signal = vec![0.0, 1.0, 0.5, 2.0, 0.0, 1.5, 0.0];
1835        let peaks = find_peaks(&signal);
1836        assert!(peaks.contains(&1));
1837        assert!(peaks.contains(&3));
1838        assert!(peaks.contains(&5));
1839    }
1840
1841    #[test]
1842    fn test_moving_average() {
1843        let signal = vec![1.0, 2.0, 3.0, 4.0, 5.0];
1844        let ma = moving_average(&signal, 3);
1845        assert_eq!(ma.len(), signal.len());
1846        assert!(approx_eq(ma[4], 4.0, EPS));
1847    }
1848
1849    #[test]
1850    fn test_snr_db() {
1851        let snr = snr_db(100.0, 1.0);
1852        assert!(approx_eq(snr, 20.0, 1e-6));
1853    }
1854
1855    #[test]
1856    fn test_generate_sine() {
1857        let sine = generate_sine(100.0, 1000.0, 100, 1.0, 0.0);
1858        assert_eq!(sine.len(), 100);
1859        // First sample should be ≈ 0
1860        assert!(sine[0].abs() < 1e-10);
1861    }
1862
1863    #[test]
1864    fn test_zero_pad() {
1865        let sig = vec![1.0, 2.0, 3.0];
1866        let padded = zero_pad(&sig, 8);
1867        assert_eq!(padded.len(), 8);
1868        assert_eq!(padded[3], 0.0);
1869    }
1870
1871    #[test]
1872    fn test_fft_frequencies() {
1873        let freqs = fft_frequencies(8, 8.0);
1874        assert_eq!(freqs.len(), 8);
1875        assert!(approx_eq(freqs[1], 1.0, EPS));
1876        assert!(approx_eq(freqs[4], 4.0, EPS));
1877    }
1878
1879    #[test]
1880    fn test_rfft_frequencies() {
1881        let freqs = rfft_frequencies(8, 8.0);
1882        assert_eq!(freqs.len(), 5); // 0..=4
1883        assert!(approx_eq(freqs[4], 4.0, EPS));
1884    }
1885
1886    #[test]
1887    fn test_real_cepstrum_length() {
1888        let signal: Vec<f64> = (0..16).map(|i| (i as f64).sin()).collect();
1889        let cep = real_cepstrum(&signal);
1890        assert_eq!(cep.len(), 16);
1891    }
1892
1893    #[test]
1894    fn test_bartlett_window_peak() {
1895        let w = bartlett_window(9);
1896        // Peak at center index 4
1897        assert!(approx_eq(w[4], 1.0, 1e-10));
1898    }
1899
1900    #[test]
1901    fn test_blackman_harris_range() {
1902        let w = blackman_harris_window(64);
1903        for &v in &w {
1904            assert!((-0.01..=1.01).contains(&v));
1905        }
1906    }
1907
1908    #[test]
1909    fn test_db2_decompose_reconstruct() {
1910        let signal = vec![1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0, 0.0];
1911        let (approx, detail) = db2_decompose(&signal);
1912        assert!(!approx.is_empty());
1913        assert!(!detail.is_empty());
1914        let recon = db2_reconstruct(&approx, &detail);
1915        assert!(!recon.is_empty());
1916    }
1917
1918    #[test]
1919    fn test_stft_istft_roundtrip_shape() {
1920        let n = 128;
1921        let signal: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
1922        let frames = stft(&signal, 32, 16, WindowType::Hann);
1923        assert!(!frames.is_empty());
1924        for frame in &frames {
1925            assert_eq!(frame.len(), 17); // 32/2 + 1
1926        }
1927    }
1928
1929    #[test]
1930    fn test_butterworth_lowpass_creates_filter() {
1931        let filt = butterworth_lowpass(4, 0.5);
1932        assert_eq!(filt.sections.len(), 2);
1933    }
1934
1935    #[test]
1936    fn test_chebyshev1_lowpass_creates_filter() {
1937        let filt = chebyshev1_lowpass(4, 0.4, 1.0);
1938        assert!(!filt.sections.is_empty());
1939    }
1940}