harmos_signal/spectrum/

fourier.rs

//! The forward transform, the averaged estimate, and the reconstruction back.

use rustfft::FftPlanner;
use rustfft::num_complex::Complex;

use super::{Band, Spectrum, Taper};

/// Transforms real samples into one single-window spectrum.
///
/// `rate` is the sampling rate in Hz, and the answer covers `0..=rate / 2` in
/// `samples.len() / 2 + 1` bins. The taper's coherent gain is divided back out
/// and every bin but DC and Nyquist is doubled, so a tone landing on a bin
/// reads its own amplitude whichever taper measured it — which is the property
/// [`reconstruct`] inverts exactly.
///
/// An empty input answers an empty spectrum; there is no error to handle.
///
/// ```
/// use harmos_signal::{Taper, transform};
///
/// let rate = 1000.0;
/// let samples: Vec<f64> = (0..1000)
///     .map(|i| 3.0 * (std::f64::consts::TAU * 50.0 * f64::from(i) / rate).sin())
///     .collect();
///
/// let spectrum = transform(&samples, rate, Taper::Rectangular);
/// let magnitude = harmos_signal::magnitude(&spectrum.real, &spectrum.imag);
///
/// // 50 Hz at a 1 Hz bin step is bin 50, and it reads the amplitude it was given.
/// assert!((magnitude[50] - 3.0).abs() < 1e-9);
/// ```
#[must_use]
pub fn transform(samples: &[f64], rate: f64, taper: Taper) -> Spectrum {
    let length = samples.len();
    if length == 0 || rate <= 0.0 {
        return Spectrum {
            step: 0.0,
            real: Vec::new(),
            imag: Vec::new(),
        };
    }

    let bins = length / 2 + 1;
    let coefficients = taper.coefficients(length);
    // The coherent gain, summed rather than averaged: dividing the transform by
    // it is what makes one taper's amplitude the same as another's.
    let gain: f64 = coefficients.iter().sum();

    let mut buffer: Vec<Complex<f64>> = samples
        .iter()
        .zip(&coefficients)
        .map(|(sample, coefficient)| Complex {
            re: sample * coefficient,
            im: 0.0,
        })
        .collect();

    FftPlanner::new()
        .plan_fft_forward(length)
        .process(&mut buffer);

    let mut real: Vec<f64> = buffer
        .iter()
        .take(bins)
        .map(|bin| bin.re / gain * 2.0)
        .collect();
    let mut imag: Vec<f64> = buffer
        .iter()
        .take(bins)
        .map(|bin| bin.im / gain * 2.0)
        .collect();

    // A single-sided spectrum doubles every bin that has a negative-frequency
    // twin. DC has none, and neither does Nyquist when the length is even.
    real[0] /= 2.0;
    imag[0] /= 2.0;
    if length.is_multiple_of(2) && bins > 1 {
        real[bins - 1] /= 2.0;
        imag[bins - 1] /= 2.0;
    }

    Spectrum {
        step: rate / length as f64,
        real,
        imag,
    }
}

/// Reconstructs real samples from a spectrum, optionally keeping one band.
///
/// The inverse of [`transform`]: the normalization it applied is undone bin by
/// bin, the negative frequencies are rebuilt as conjugates, and the result is
/// the sample run the spectrum came from. Passing a [`Band`] scales each bin by
/// that band's coefficient first, which is how a band-limited copy of a signal
/// is produced without designing a filter.
///
/// The answer holds `(bins - 1) * 2` samples at the rate [`Spectrum::rate`]
/// reports. A spectrum of fewer than two bins reconstructs to nothing.
#[must_use]
pub fn reconstruct(spectrum: &Spectrum, band: Option<Band>) -> Vec<f64> {
    let bins = spectrum.real.len();
    if bins < 2 {
        return Vec::new();
    }

    let length = (bins - 1) * 2;
    let mut buffer: Vec<Complex<f64>> = vec![Complex::default(); length];

    for (index, (real, imag)) in spectrum.real.iter().zip(&spectrum.imag).enumerate() {
        let coefficient = band.map_or(1.0, |band| band.coefficient(index as f64 * spectrum.step));
        // Undo exactly what the forward transform did: DC and Nyquist were
        // never doubled, so they are scaled by the whole length and the rest by
        // half of it.
        let scale = if index == 0 || index == bins - 1 {
            length as f64
        } else {
            length as f64 / 2.0
        };

        buffer[index] = Complex {
            re: real * coefficient * scale,
            im: imag * coefficient * scale,
        };
    }

    for index in 1..bins - 1 {
        buffer[length - index] = buffer[index].conj();
    }

    FftPlanner::new()
        .plan_fft_inverse(length)
        .process(&mut buffer);

    buffer
        .iter()
        .map(|sample| sample.re / length as f64)
        .collect()
}

/// Welch's method: what resolution to estimate at, and how hard to average.
///
/// Averaging the periodograms of overlapping segments trades frequency
/// resolution for a quieter estimate, which is the trade a noisy measurement
/// wants and a phase-sensitive one does not. [`Welch::estimate`] is the run.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Welch {
    /// The frequency resolution to aim for, in Hz.
    ///
    /// One segment is `rate / resolution` samples long, so a finer resolution
    /// buys more bins out of fewer averages.
    pub resolution: f64,
    /// How much of each segment the next one repeats, `0.0..1.0`.
    pub overlap: f64,
    /// The taper applied to every segment.
    pub taper: Taper,
}

impl Default for Welch {
    fn default() -> Self {
        Self {
            resolution: 100.0,
            overlap: 0.5,
            taper: Taper::Hann,
        }
    }
}

impl Welch {
    /// A default estimate at a stated resolution: half overlap, Hann taper.
    #[must_use]
    pub fn new(resolution: f64) -> Self {
        Self {
            resolution,
            ..Self::default()
        }
    }

    /// Estimates the amplitude spectrum of `samples` sampled at `rate` Hz.
    ///
    /// Segments shorter than 64 samples are not worth averaging and the segment
    /// length is clamped there; a signal too short to divide is transformed
    /// once instead, so an estimate is always answered on the same terms.
    ///
    /// Phase is not part of the answer, and that is the method rather than an
    /// omission: averaging periodograms discards it.
    #[must_use]
    pub fn estimate(&self, samples: &[f64], rate: f64) -> Estimate {
        let total = samples.len();
        if total == 0 || rate <= 0.0 {
            return Estimate {
                step: 0.0,
                magnitude: Vec::new(),
            };
        }

        let segment = ((rate / self.resolution).round() as usize).clamp(64, total);
        let hop =
            (((segment as f64) * (1.0 - self.overlap.clamp(0.0, 0.99))).round() as usize).max(1);
        let segments = if total <= segment {
            1
        } else {
            (total - segment) / hop + 1
        };

        if segments <= 1 || segment >= total {
            let single = transform(samples, rate, self.taper);
            return Estimate {
                step: single.step,
                magnitude: super::magnitude(&single.real, &single.imag),
            };
        }

        let bins = segment / 2 + 1;
        let coefficients = self.taper.coefficients(segment);
        let gain: f64 = coefficients.iter().sum();

        let mut power = vec![0.0; bins];
        let mut planner = FftPlanner::new();
        let plan = planner.plan_fft_forward(segment);

        for index in 0..segments {
            let start = index * hop;
            let Some(window) = samples.get(start..start + segment) else {
                break;
            };

            let mut buffer: Vec<Complex<f64>> = window
                .iter()
                .zip(&coefficients)
                .map(|(sample, coefficient)| Complex {
                    re: sample * coefficient,
                    im: 0.0,
                })
                .collect();

            plan.process(&mut buffer);

            for (accumulated, bin) in power.iter_mut().zip(buffer.iter().take(bins)) {
                *accumulated += bin.norm_sqr();
            }
        }

        // Summed over K segments the power of a tone is K·A²·(Σw)²/4, so
        // dividing by K·(Σw)² and doubling the root recovers A.
        let normalizer = segments as f64 * gain.powi(2);
        let mut magnitude: Vec<f64> = power
            .iter()
            .map(|power| (power / normalizer).sqrt() * 2.0)
            .collect();

        magnitude[0] /= 2.0;
        if segment.is_multiple_of(2) && bins > 1 {
            magnitude[bins - 1] /= 2.0;
        }

        Estimate {
            step: rate / segment as f64,
            magnitude,
        }
    }
}

/// One averaged spectral estimate: a bin step, and a magnitude per bin.
///
/// Deliberately not a [`Spectrum`]. Welch's method has no phase to report, and
/// a spectrum carrying an invented zero imaginary part would reconstruct to
/// something that was never measured.
#[derive(Clone, Debug, PartialEq)]
pub struct Estimate {
    /// The frequency distance between neighbouring bins, in Hz.
    pub step: f64,
    /// The amplitude of each bin, from DC upwards.
    pub magnitude: Vec<f64>,
}

#[cfg(test)]
mod tests {
    use std::f64::consts::TAU;

    use super::*;
    use crate::spectrum::{bins, magnitude, phase};

    /// A tone of `amplitude` at `frequency` Hz, `count` samples at `rate` Hz.
    fn tone(amplitude: f64, frequency: f64, rate: f64, count: usize) -> Vec<f64> {
        (0..count)
            .map(|i| amplitude * (TAU * frequency * i as f64 / rate).sin())
            .collect()
    }

    #[test]
    fn a_constant_signal_is_all_dc() {
        let spectrum = transform(&[5.0; 1000], 1000.0, Taper::Rectangular);
        let magnitude = magnitude(&spectrum.real, &spectrum.imag);

        assert!((magnitude[0] - 5.0).abs() < 1e-9);
        assert!(magnitude[1..].iter().all(|bin| *bin < 1e-9));
    }

    #[test]
    fn an_on_bin_tone_reads_its_own_amplitude_under_every_taper() {
        let samples = tone(10.0, 50.0, 1000.0, 1000);

        for taper in [
            Taper::Rectangular,
            Taper::Hann,
            Taper::Hamming,
            Taper::Blackman,
        ] {
            let spectrum = transform(&samples, 1000.0, taper);
            let peak = magnitude(&spectrum.real, &spectrum.imag)
                .into_iter()
                .fold(0.0_f64, f64::max);

            assert!((peak - 10.0).abs() < 0.15, "{taper:?} recovered {peak}");
        }
    }

    #[test]
    fn a_sine_lags_a_cosine_by_a_quarter_turn() {
        let rate = 1000.0;
        let sine = transform(&tone(1.0, 100.0, rate, 1000), rate, Taper::Rectangular);
        let cosine: Vec<f64> = (0..1000)
            .map(|i| (TAU * 100.0 * f64::from(i) / rate).cos())
            .collect();
        let cosine = transform(&cosine, rate, Taper::Rectangular);

        assert!((phase(&sine.real, &sine.imag)[100] + std::f64::consts::FRAC_PI_2).abs() < 1e-6);
        assert!(phase(&cosine.real, &cosine.imag)[100].abs() < 1e-6);
    }

    #[test]
    fn the_spectrum_accounts_for_every_bit_of_the_signal() {
        // Parseval, as this crate's normalization leaves it. A single-sided
        // spectrum doubled every bin with a negative-frequency twin, so folding
        // the power back counts those half and DC and Nyquist whole. Three
        // things have to agree for this to hold: the coherent-gain division,
        // the DC and Nyquist half-scaling, and the time-domain fold.
        let rate = 1000.0;
        let samples: Vec<f64> = (0..1000)
            .map(|i| {
                let t = f64::from(i) / rate;
                3.0 + 8.0 * (TAU * 25.0 * t).cos() + 4.0 * (TAU * 75.0 * t).sin()
            })
            .collect();

        let spectrum = transform(&samples, rate, Taper::Rectangular);
        let amplitude = magnitude(&spectrum.real, &spectrum.imag);
        let nyquist = amplitude.len() - 1;

        let folded: f64 = amplitude[0].powi(2)
            + amplitude[nyquist].powi(2)
            + amplitude[1..nyquist]
                .iter()
                .map(|bin| bin * bin)
                .sum::<f64>()
                / 2.0;

        // 3² + 8²/2 + 4²/2 = 49, so the signal's rms is exactly seven.
        let measured = crate::waveform::Summary::of(&samples)
            .rms()
            .expect("a thousand samples have an rms");

        assert!(
            (folded - 49.0).abs() < 1e-9,
            "the spectrum accounted for {folded}"
        );
        assert!(
            (measured - 7.0).abs() < 1e-9,
            "the samples measured {measured}"
        );
        assert!(
            (folded.sqrt() - measured).abs() < 1e-9,
            "the two domains disagree"
        );
    }

    #[test]
    fn the_bin_step_and_the_rate_are_each_other() {
        let spectrum = transform(&tone(1.0, 50.0, 1000.0, 1000), 1000.0, Taper::Hann);

        assert_eq!(spectrum.real.len(), 501);
        assert!((spectrum.step - 1.0).abs() < 1e-12);
        assert!((spectrum.rate() - 1000.0).abs() < 1e-9);
        assert!((bins(spectrum.step, spectrum.real.len())[500] - 500.0).abs() < 1e-9);
    }

    #[test]
    fn a_taper_buys_lower_sidelobes_off_bin() {
        // 47.5 Hz sits between bins, so the leakage a taper suppresses is there
        // to suppress.
        let samples = tone(10.0, 47.5, 1000.0, 1000);
        let skirt = |taper| {
            let spectrum = transform(&samples, 1000.0, taper);
            magnitude(&spectrum.real, &spectrum.imag)
                .into_iter()
                .enumerate()
                .filter(|(bin, _)| bin.abs_diff(48) > 5)
                .map(|(_, value)| value)
                .sum::<f64>()
        };

        assert!(skirt(Taper::Hann) < skirt(Taper::Rectangular));
    }

    #[test]
    fn an_empty_run_transforms_to_an_empty_spectrum() {
        let spectrum = transform(&[], 1000.0, Taper::Hann);

        assert!(spectrum.real.is_empty());
        assert!(reconstruct(&spectrum, None).is_empty());
    }

    #[test]
    fn a_transform_and_its_reconstruction_return_the_samples() {
        let rate = 1000.0;
        let samples: Vec<f64> = (0..1000)
            .map(|i| {
                let t = f64::from(i) / rate;
                3.0 + 8.0 * (TAU * 25.0 * t).cos() + 4.0 * (TAU * 75.0 * t).sin()
            })
            .collect();

        let recovered = reconstruct(&transform(&samples, rate, Taper::Rectangular), None);

        assert_eq!(recovered.len(), samples.len());
        for (original, recovered) in samples.iter().zip(&recovered) {
            assert!((original - recovered).abs() < 1e-9);
        }
    }

    #[test]
    fn a_band_keeps_the_component_inside_it_and_drops_the_rest() {
        let rate = 1000.0;
        let samples: Vec<f64> = (0..1000)
            .map(|i| {
                let t = f64::from(i) / rate;
                5.0 + 10.0 * (TAU * 50.0 * t).sin() + 3.0 * (TAU * 150.0 * t).sin()
            })
            .collect();

        let spectrum = transform(&samples, rate, Taper::Rectangular);
        let band = Band::Rectangular {
            start: 40.0,
            end: 60.0,
        };
        let filtered = reconstruct(&spectrum, Some(band));

        // What survives is the 50 Hz tone alone, whose RMS is its amplitude
        // over root two.
        let power: f64 = filtered.iter().map(|sample| sample * sample).sum();
        let rms = (power / filtered.len() as f64).sqrt();

        assert!(
            (rms - 10.0 / 2.0_f64.sqrt()).abs() < 0.05,
            "reconstructed rms {rms}"
        );
    }

    #[test]
    fn welch_and_a_single_window_agree_on_a_tone() {
        let rate = 10_000.0;
        let samples = tone(10.0, 100.0, rate, 10_000);

        let single = transform(&samples, rate, Taper::Hann);
        let single = magnitude(&single.real, &single.imag)
            .into_iter()
            .fold(0.0_f64, f64::max);
        let averaged = Welch::new(50.0)
            .estimate(&samples, rate)
            .magnitude
            .into_iter()
            .fold(0.0_f64, f64::max);

        assert!((single - 10.0).abs() < 2.0, "single window read {single}");
        assert!((averaged - 10.0).abs() < 2.0, "welch read {averaged}");
    }

    #[test]
    fn welch_finds_a_tone_a_single_window_would_bury() {
        let rate = 10_000.0;
        let samples: Vec<f64> = (0..10_000)
            .map(|i| {
                let noise = f64::from((i * 12_345 + 67_890) % 1000) / 1000.0 - 0.5;
                10.0 * (TAU * 100.0 * f64::from(i) / rate).sin() + noise
            })
            .collect();

        let estimate = Welch::default().estimate(&samples, rate);
        let peak = estimate
            .magnitude
            .iter()
            .enumerate()
            .max_by(|(_, left), (_, right)| left.total_cmp(right))
            .map(|(bin, _)| bin as f64 * estimate.step)
            .expect("a non-empty estimate has a peak");

        assert!((peak - 100.0).abs() < 5.0, "welch peaked at {peak} Hz");
    }

    #[test]
    fn a_run_too_short_to_average_is_transformed_once() {
        let rate = 1000.0;
        let samples = tone(1.0, 50.0, rate, 200);

        let estimate = Welch::new(1.0).estimate(&samples, rate);
        let single = transform(&samples, rate, Taper::Hann);

        assert_eq!(estimate.magnitude, magnitude(&single.real, &single.imag));
        assert!((estimate.step - single.step).abs() < 1e-12);
    }

    #[test]
    fn an_empty_run_estimates_to_nothing() {
        let estimate = Welch::default().estimate(&[], 1000.0);

        assert!(estimate.magnitude.is_empty());
        assert_eq!(estimate.step, 0.0);
    }
}