rosa 0.1.0

Audio feature extraction library matching librosa output
Documentation
//! Tempo estimation matching `librosa.feature.rhythm.tempo`.

use crate::matrix::Matrix;
use crate::stft::fft_frequencies;

/// Map autocorrelation lag bins to BPM.
///
/// Matches `librosa.tempo_frequencies`.
/// `bin_frequencies[0] = +inf`, `bin_frequencies[i] = 60 * sr / (hop_length * i)`.
pub fn tempo_frequencies(n_bins: usize, hop_length: usize, sr: f64) -> Vec<f64> {
    let mut freqs = vec![0.0; n_bins];
    if n_bins > 0 {
        freqs[0] = f64::INFINITY;
    }
    for i in 1..n_bins {
        freqs[i] = 60.0 * sr / (hop_length as f64 * i as f64);
    }
    freqs
}

/// Map Fourier tempogram frequency bins to BPM.
///
/// Matches `librosa.fourier_tempo_frequencies`.
pub fn fourier_tempo_frequencies(sr: f64, win_length: usize, hop_length: usize) -> Vec<f64> {
    let sr_bpm = sr * 60.0 / hop_length as f64;
    fft_frequencies(sr_bpm, win_length)
}

/// Parameters for tempo estimation.
pub struct TempoParams {
    pub sr: f64,
    pub hop_length: usize,
    pub start_bpm: f64,
    pub std_bpm: f64,
    pub ac_size: f64,
    pub max_tempo: Option<f64>,
}

impl TempoParams {
    pub fn new(sr: f64) -> Self {
        Self {
            sr,
            hop_length: 512,
            start_bpm: 120.0,
            std_bpm: 1.0,
            ac_size: 8.0,
            max_tempo: Some(320.0),
        }
    }
}

/// Estimate global tempo in BPM from onset envelope.
///
/// Matches `librosa.feature.tempo` with `aggregate=np.mean`.
pub fn tempo(onset_envelope: &[f64], params: &TempoParams) -> f64 {
    // Compute win_length from ac_size
    let win_length = (params.ac_size * params.sr / params.hop_length as f64).round() as usize;

    if win_length < 2 {
        return params.start_bpm;
    }

    // Compute tempogram
    let tg_params = super::tempogram::TempogramParams {
        sr: params.sr,
        hop_length: params.hop_length,
        win_length,
        center: true,
        norm: Some(crate::dsp::util::Norm::Linf),
    };
    let tg = super::tempogram::tempogram(onset_envelope, &tg_params);

    tempo_from_tempogram(&tg, win_length, params)
}

/// Estimate tempo from a pre-computed tempogram.
pub(crate) fn tempo_from_tempogram(tg: &Matrix, win_length: usize, params: &TempoParams) -> f64 {
    // Aggregate: mean across time axis → Vec of length win_length
    let n_frames = tg.cols();
    let mut tg_mean = vec![0.0; win_length];
    if n_frames > 0 {
        for lag in 0..win_length {
            let mut sum = 0.0;
            for t in 0..n_frames {
                sum += tg.get(lag, t);
            }
            tg_mean[lag] = sum / n_frames as f64;
        }
    }

    // Get BPM bins
    let bpms = tempo_frequencies(win_length, params.hop_length, params.sr);

    // Compute log-normal prior
    let mut logprior = vec![f64::NEG_INFINITY; win_length];
    let log2_start = params.start_bpm.log2();
    for i in 1..win_length {
        if bpms[i].is_finite() && bpms[i] > 0.0 {
            let log2_bpm = bpms[i].log2();
            logprior[i] = -0.5 * ((log2_bpm - log2_start) / params.std_bpm).powi(2);
        }
    }

    // Zero out bins above max_tempo
    if let Some(max_tempo) = params.max_tempo {
        // bpms is in descending order (bin 0=inf, bin 1=highest)
        // Find first bin where bpm < max_tempo
        for i in 0..win_length {
            if bpms[i] < max_tempo {
                break;
            }
            logprior[i] = f64::NEG_INFINITY;
        }
    }

    // Find best tempo: argmax(log1p(1e6 * tg) + logprior)
    let mut best_idx = 1; // skip bin 0
    let mut best_score = f64::NEG_INFINITY;
    for i in 1..win_length {
        let score = (1e6_f64 * tg_mean[i]).ln_1p() + logprior[i];
        if score > best_score {
            best_score = score;
            best_idx = i;
        }
    }

    bpms[best_idx]
}