oximedia-audio-analysis 0.1.7

//! Onset detection for audio analysis.
//!
//! This module provides multiple methods for detecting the onset of musical notes
//! or percussive events in audio signals.

/// Method used for onset detection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OnsetMethod {
    /// Energy-based onset detection using frame energy deltas.
    EnergyBased,
    /// Spectral flux: sum of positive spectral differences between frames.
    SpectralFlux,
    /// Phase deviation: deviation from expected phase advancement.
    PhaseDeviation,
    /// Complex domain: combines magnitude and phase information.
    ComplexDomain,
}

impl OnsetMethod {
    /// Typical latency introduced by this detection method, in milliseconds.
    #[must_use]
    pub fn typical_latency_ms(&self) -> f32 {
        match self {
            Self::EnergyBased => 10.0,
            Self::SpectralFlux | Self::PhaseDeviation | Self::ComplexDomain => 23.2,
        }
    }
}

/// An onset event detected in the audio signal.
#[derive(Debug, Clone)]
pub struct OnsetEvent {
    /// Time of onset in milliseconds from the start of the signal.
    pub time_ms: u64,
    /// Strength of the onset (0.0 = weak, 1.0 = strong).
    pub strength: f32,
    /// Detection method that produced this event.
    pub method: OnsetMethod,
}

impl OnsetEvent {
    /// Returns `true` if this onset's strength exceeds the given threshold.
    #[must_use]
    pub fn is_strong(&self, threshold: f32) -> bool {
        self.strength >= threshold
    }
}

/// Compute the spectral flux between two consecutive spectra.
///
/// Spectral flux is the sum of positive differences in magnitude between
/// the current and previous spectra, measuring how much new energy appears.
///
/// # Panics
/// Panics if `prev_spectrum` and `curr_spectrum` have different lengths.
#[must_use]
pub fn spectral_flux(prev_spectrum: &[f32], curr_spectrum: &[f32]) -> f32 {
    assert_eq!(
        prev_spectrum.len(),
        curr_spectrum.len(),
        "Spectra must have the same length"
    );
    curr_spectrum
        .iter()
        .zip(prev_spectrum.iter())
        .map(|(&curr, &prev)| {
            let diff = curr - prev;
            if diff > 0.0 {
                diff
            } else {
                0.0
            }
        })
        .sum()
}

/// Compute per-frame energy deltas from a flat sample buffer.
///
/// Each frame of `hop_size` samples has its RMS energy computed. The output
/// is the absolute difference in energy between successive frames.
#[must_use]
pub fn energy_onset(frames: &[f32], hop_size: usize) -> Vec<f32> {
    if hop_size == 0 || frames.is_empty() {
        return Vec::new();
    }

    let frame_energies: Vec<f32> = frames
        .chunks(hop_size)
        .map(|chunk| {
            let sum_sq: f32 = chunk.iter().map(|&s| s * s).sum();
            (sum_sq / chunk.len() as f32).sqrt()
        })
        .collect();

    if frame_energies.len() < 2 {
        return Vec::new();
    }

    frame_energies
        .windows(2)
        .map(|w| (w[1] - w[0]).abs())
        .collect()
}

/// Picks local peaks in a novelty curve above `threshold`.
fn peak_pick(novelty: &[f32], threshold: f32) -> Vec<usize> {
    if novelty.len() < 3 {
        return Vec::new();
    }
    let mut peaks = Vec::new();
    for i in 1..novelty.len() - 1 {
        if novelty[i] > threshold && novelty[i] >= novelty[i - 1] && novelty[i] >= novelty[i + 1] {
            peaks.push(i);
        }
    }
    peaks
}

/// Onset detector that analyses a signal and returns detected onset events.
pub struct OnsetDetector {
    /// Detection method.
    pub method: OnsetMethod,
    /// Threshold for novelty curve peak picking (0.0 – 1.0).
    pub threshold: f32,
    /// Sample rate of the audio signal.
    pub sample_rate: u32,
}

impl OnsetDetector {
    /// Create a new `OnsetDetector` with the given parameters.
    #[must_use]
    pub fn new(method: OnsetMethod, threshold: f32, sample_rate: u32) -> Self {
        Self {
            method,
            threshold,
            sample_rate,
        }
    }

    /// Detect onsets in the given audio signal.
    ///
    /// Returns a list of [`OnsetEvent`]s sorted by time.
    #[must_use]
    pub fn detect(&self, signal: &[f32]) -> Vec<OnsetEvent> {
        let hop_size = (self.sample_rate / 100) as usize; // 10 ms hops
        let hop_size = hop_size.max(1);

        let novelty = match self.method {
            OnsetMethod::EnergyBased => energy_onset(signal, hop_size),
            OnsetMethod::SpectralFlux
            | OnsetMethod::PhaseDeviation
            | OnsetMethod::ComplexDomain => {
                // Use energy as a proxy when no FFT buffer is available
                energy_onset(signal, hop_size)
            }
        };

        if novelty.is_empty() {
            return Vec::new();
        }

        // Normalise novelty to [0, 1]
        let max_val = novelty.iter().copied().fold(0.0_f32, f32::max);
        let normalised: Vec<f32> = if max_val > 0.0 {
            novelty.iter().map(|&v| v / max_val).collect()
        } else {
            novelty.clone()
        };

        let peaks = peak_pick(&normalised, self.threshold);

        peaks
            .into_iter()
            .map(|frame_idx| {
                let time_ms =
                    (frame_idx as u64 * hop_size as u64 * 1000) / u64::from(self.sample_rate);
                OnsetEvent {
                    time_ms,
                    strength: normalised[frame_idx],
                    method: self.method,
                }
            })
            .collect()
    }
}

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

    #[test]
    fn test_onset_method_latency_energy() {
        assert!((OnsetMethod::EnergyBased.typical_latency_ms() - 10.0).abs() < 0.01);
    }

    #[test]
    fn test_onset_method_latency_flux() {
        assert!(OnsetMethod::SpectralFlux.typical_latency_ms() > 0.0);
    }

    #[test]
    fn test_onset_method_latency_phase() {
        assert!(OnsetMethod::PhaseDeviation.typical_latency_ms() > 0.0);
    }

    #[test]
    fn test_onset_method_latency_complex() {
        assert!(OnsetMethod::ComplexDomain.typical_latency_ms() > 0.0);
    }

    #[test]
    fn test_onset_event_is_strong() {
        let ev = OnsetEvent {
            time_ms: 100,
            strength: 0.8,
            method: OnsetMethod::EnergyBased,
        };
        assert!(ev.is_strong(0.5));
        assert!(!ev.is_strong(0.9));
    }

    #[test]
    fn test_onset_event_threshold_boundary() {
        let ev = OnsetEvent {
            time_ms: 0,
            strength: 0.5,
            method: OnsetMethod::SpectralFlux,
        };
        assert!(ev.is_strong(0.5));
    }

    #[test]
    fn test_spectral_flux_positive_only() {
        let prev = vec![1.0, 2.0, 3.0];
        let curr = vec![2.0, 1.0, 5.0];
        // Positive differences: (2-1)=1.0 and (5-3)=2.0; negative (1-2) ignored
        let flux = spectral_flux(&prev, &curr);
        assert!((flux - 3.0).abs() < 1e-5);
    }

    #[test]
    fn test_spectral_flux_no_increase() {
        let prev = vec![5.0, 5.0, 5.0];
        let curr = vec![1.0, 2.0, 3.0];
        let flux = spectral_flux(&prev, &curr);
        assert_eq!(flux, 0.0);
    }

    #[test]
    fn test_spectral_flux_equal_spectra() {
        let spec = vec![1.0, 2.0, 3.0];
        assert_eq!(spectral_flux(&spec, &spec), 0.0);
    }

    #[test]
    fn test_energy_onset_empty() {
        assert!(energy_onset(&[], 512).is_empty());
    }

    #[test]
    fn test_energy_onset_returns_deltas() {
        // Two frames: silent then loud
        let mut signal = vec![0.0_f32; 512];
        signal.extend(vec![1.0_f32; 512]);
        let deltas = energy_onset(&signal, 512);
        assert!(!deltas.is_empty());
        assert!(deltas[0] > 0.0); // Energy jumped
    }

    #[test]
    fn test_onset_detector_silent_signal() {
        let detector = OnsetDetector::new(OnsetMethod::EnergyBased, 0.5, 44100);
        let signal = vec![0.0_f32; 44100];
        let events = detector.detect(&signal);
        assert!(events.is_empty());
    }

    #[test]
    fn test_onset_detector_finds_attack() {
        let detector = OnsetDetector::new(OnsetMethod::EnergyBased, 0.5, 44100);
        // Impulse at sample 22050 (0.5 s)
        let mut signal = vec![0.0_f32; 44100];
        for s in &mut signal[22050..22060] {
            *s = 1.0;
        }
        let events = detector.detect(&signal);
        // There should be at least one onset detected
        assert!(!events.is_empty());
    }

    #[test]
    fn test_onset_detector_spectral_flux_method() {
        let detector = OnsetDetector::new(OnsetMethod::SpectralFlux, 0.3, 22050);
        let mut signal = vec![0.0_f32; 22050];
        for s in &mut signal[11000..11020] {
            *s = 0.9;
        }
        let events = detector.detect(&signal);
        // Should produce some events given the impulse
        let _ = events; // Just verify it doesn't panic
    }
}