rosa 0.1.0

Audio feature extraction library matching librosa output
Documentation
//! Onset detection and backtracking matching `librosa.onset.onset_detect`.

use crate::dsp::util::peak_pick;

/// Detect onset events by peak picking the onset envelope.
///
/// Matches `librosa.onset.onset_detect(onset_envelope=..., sr=sr, hop_length=hop_length, units='frames')`.
///
/// - `onset_envelope`: pre-computed onset strength envelope
/// - `sr`: sample rate
/// - `hop_length`: hop length used to compute the envelope
/// - `normalize`: if true, normalize envelope to [0, 1] before peak picking (default: true)
/// - `backtrack`: if true, backtrack onsets to preceding local minima
///
/// # Examples
///
/// ```
/// let envelope = vec![0.0, 0.1, 0.5, 0.8, 0.3, 0.1, 0.0, 0.2, 0.9, 0.4];
/// let onsets = rosa::onset::onset_detect(&envelope, 22050.0, 512, true, false);
/// // Returns frame indices of detected onsets
/// ```
pub fn onset_detect(
    onset_envelope: &[f64],
    sr: f64,
    hop_length: usize,
    normalize: bool,
    backtrack: bool,
) -> Vec<usize> {
    if onset_envelope.is_empty() {
        return Vec::new();
    }

    // Normalize to [0, 1]
    let env = if normalize {
        let min_val = onset_envelope.iter().copied().fold(f64::INFINITY, f64::min);
        let shifted: Vec<f64> = onset_envelope.iter().map(|&v| v - min_val).collect();
        let max_val = shifted.iter().copied().fold(0.0_f64, f64::max);
        let tiny = (f64::EPSILON * max_val.abs()).sqrt().max(f64::MIN_POSITIVE);
        shifted
            .iter()
            .map(|&v| v / (max_val + tiny))
            .collect::<Vec<_>>()
    } else {
        onset_envelope.to_vec()
    };

    // Check if any onsets are present
    if !env.iter().any(|&v| v != 0.0) || !env.iter().all(|v| v.is_finite()) {
        return Vec::new();
    }

    // Peak pick parameters tuned by large-scale search (librosa defaults)
    let pre_max = ((0.03 * sr / hop_length as f64).floor() as usize).max(1);
    let post_max = ((0.00 * sr / hop_length as f64).floor() as usize) + 1;
    let pre_avg = ((0.10 * sr / hop_length as f64).floor() as usize).max(1);
    let post_avg = ((0.10 * sr / hop_length as f64).floor() as usize) + 1;
    let wait = ((0.03 * sr / hop_length as f64).floor() as usize).max(1);
    let delta = 0.07;

    let mut onsets = peak_pick(&env, pre_max, post_max, pre_avg, post_avg, delta, wait);

    // Optional backtracking
    if backtrack {
        onsets = onset_backtrack(&onsets, &env);
    }

    onsets
}

/// Backtrack detected onset events to the nearest preceding local minimum.
///
/// Matches `librosa.onset.onset_backtrack(events, energy)`.
///
/// For each detected onset frame, finds the nearest preceding local minimum
/// in the energy signal.
pub fn onset_backtrack(events: &[usize], energy: &[f64]) -> Vec<usize> {
    if events.is_empty() || energy.is_empty() {
        return events.to_vec();
    }

    // Find local minima in energy
    // librosa: (energy[1:-1] <= energy[:-2]) & (energy[1:-1] < energy[2:])
    // This is different from our localmin which uses strict < on left.
    // librosa backtrack uses: <=left AND <right
    let n = energy.len();
    let mut minima_mask = vec![false; n];
    for i in 1..n.saturating_sub(1) {
        minima_mask[i] = energy[i] <= energy[i - 1] && energy[i] < energy[i + 1];
    }

    // Collect minima indices (shift by 1 to account for slicing, already correct)
    let mut minima: Vec<usize> = minima_mask
        .iter()
        .enumerate()
        .filter(|&(_, &b)| b)
        .map(|(i, _)| i)
        .collect();

    // Ensure frame 0 is included (fix_frames with x_min=0)
    if minima.is_empty() || minima[0] != 0 {
        minima.insert(0, 0);
    }

    // Match each event to nearest preceding minimum (right=False in librosa)
    events
        .iter()
        .map(|&event| {
            // Binary search: find the rightmost minimum <= event
            match minima.binary_search(&event) {
                Ok(idx) => minima[idx],
                Err(idx) => {
                    if idx == 0 {
                        minima[0]
                    } else {
                        minima[idx - 1]
                    }
                }
            }
        })
        .collect()
}