peerspeak/src/dsp/metrics.rs

//! Objective signal metrics — the numbers that turn "sounds better" into a
//! measurement. The headline one for echo cancellation is **ERLE** (Echo Return
//! Loss Enhancement): how much echo energy the canceller removed, in dB.

/// Root-mean-square level of a signal (linear amplitude). `0.0` for empty input.
pub fn rms(samples: &[f32]) -> f32 {
    if samples.is_empty() {
        return 0.0;
    }
    let sum_sq: f64 = samples.iter().map(|&s| s as f64 * s as f64).sum();
    (sum_sq / samples.len() as f64).sqrt() as f32
}

/// Peak absolute amplitude. `0.0` for empty input.
pub fn peak(samples: &[f32]) -> f32 {
    samples.iter().fold(0.0, |m, &s| m.max(s.abs()))
}

/// Converts a linear amplitude (e.g. an RMS or peak value, relative to full-scale
/// `1.0`) to decibels below full scale. A floor of -120 dBFS is returned for
/// silence so the result is always finite.
pub fn dbfs(linear: f32) -> f32 {
    if linear <= 1e-6 {
        return -120.0;
    }
    20.0 * linear.log10()
}

/// **Echo Return Loss Enhancement**, in dB: `10·log10(E[echo²] / E[residual²])`.
///
/// `echo` is the signal *before* cancellation (the echo the mic picked up);
/// `residual` is what's *left after* the canceller ran. A larger number is
/// better — e.g. +30 dB means the canceller removed 99.9% of the echo energy.
/// Returns a +120 dB ceiling if the residual is effectively silent (perfect
/// cancellation) and 0.0 if there was no echo energy to begin with.
pub fn erle(echo: &[f32], residual: &[f32]) -> f32 {
    let e_echo = mean_square(echo);
    let e_res = mean_square(residual);
    if e_echo <= 1e-12 {
        return 0.0;
    }
    if e_res <= 1e-12 {
        return 120.0;
    }
    10.0 * (e_echo / e_res).log10() as f32
}

/// Mean square (average energy per sample) of a signal. `0.0` for empty input.
fn mean_square(samples: &[f32]) -> f64 {
    if samples.is_empty() {
        return 0.0;
    }
    samples.iter().map(|&s| s as f64 * s as f64).sum::<f64>() / samples.len() as f64
}

/// A frequency band for per-band energy analysis, in Hz.
pub struct Band {
    pub label: &'static str,
    pub low_hz: f32,
    pub high_hz: f32,
}

/// Voice-relevant bands for spotting *where* residual echo or noise lives.
pub const VOICE_BANDS: &[Band] = &[
    Band {
        label: "low (80-300)",
        low_hz: 80.0,
        high_hz: 300.0,
    },
    Band {
        label: "low-mid (300-1k)",
        low_hz: 300.0,
        high_hz: 1000.0,
    },
    Band {
        label: "mid (1k-3k)",
        low_hz: 1000.0,
        high_hz: 3000.0,
    },
    Band {
        label: "high-mid (3k-6k)",
        low_hz: 3000.0,
        high_hz: 6000.0,
    },
    Band {
        label: "high (6k-12k)",
        low_hz: 6000.0,
        high_hz: 12000.0,
    },
];

/// Sums the linear magnitude energy within `[low_hz, high_hz)` across a single
/// STFT magnitude frame. `bin_hz` maps a bin index to its centre frequency.
pub fn band_energy(frame: &[f32], bin_hz: impl Fn(usize) -> f32, low_hz: f32, high_hz: f32) -> f32 {
    frame
        .iter()
        .enumerate()
        .filter(|&(bin, _)| {
            let hz = bin_hz(bin);
            hz >= low_hz && hz < high_hz
        })
        .map(|(_, &m)| m * m)
        .sum()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rms_of_constant_is_that_constant() {
        assert!((rms(&[0.5; 100]) - 0.5).abs() < 1e-6);
        assert_eq!(rms(&[]), 0.0);
    }

    #[test]
    fn peak_finds_max_magnitude() {
        assert_eq!(peak(&[0.1, -0.9, 0.3]), 0.9);
    }

    #[test]
    fn dbfs_landmarks() {
        assert!((dbfs(1.0) - 0.0).abs() < 1e-4, "full scale = 0 dBFS");
        assert!((dbfs(0.5) - -6.0206).abs() < 1e-3, "half = ~-6 dB");
        assert_eq!(dbfs(0.0), -120.0, "silence floors");
    }

    #[test]
    fn erle_halving_energy_is_about_3db() {
        // residual amplitude = echo/sqrt(2) -> half the energy -> ~3.01 dB.
        let echo = vec![1.0f32; 1000];
        let residual = vec![std::f32::consts::FRAC_1_SQRT_2; 1000];
        let e = erle(&echo, &residual);
        assert!((e - 3.0103).abs() < 0.01, "expected ~3 dB, got {e}");
    }

    #[test]
    fn erle_perfect_cancellation_ceils() {
        assert_eq!(erle(&[1.0; 10], &[0.0; 10]), 120.0);
        assert_eq!(erle(&[0.0; 10], &[0.0; 10]), 0.0, "no echo -> 0");
    }

    #[test]
    fn band_energy_selects_the_right_bins() {
        // 5-bin frame, 100 Hz per bin: bins at 0,100,200,300,400 Hz.
        let frame = [1.0, 2.0, 3.0, 4.0, 5.0];
        let bin_hz = |b: usize| b as f32 * 100.0;
        // [150,350) -> bins 200,300 Hz -> 3^2 + 4^2 = 25.
        let e = band_energy(&frame, bin_hz, 150.0, 350.0);
        assert!((e - 25.0).abs() < 1e-4, "got {e}");
    }
}