rosa 0.1.0

Audio feature extraction library matching librosa output
Documentation
use rosa::beat::*;

fn assert_approx_eq(actual: &[f64], expected: &[f64], tol: f64, name: &str) {
    assert_eq!(
        actual.len(),
        expected.len(),
        "{name}: length mismatch: got {} expected {}",
        actual.len(),
        expected.len()
    );
    for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() {
        // Handle infinity (JSON encodes inf as 1e308)
        if e.is_infinite() || *e >= 1e307 {
            assert!(
                a.is_infinite() || *a >= 1e307,
                "{name}[{i}]: expected inf/very large, got {a}"
            );
            continue;
        }
        let diff = (a - e).abs();
        let threshold = tol * e.abs().max(1.0);
        assert!(
            diff < threshold,
            "{name}[{i}]: got {a}, expected {e}, diff {diff} > tol {threshold}"
        );
    }
}

// ---------------------------------------------------------------------------
// pad_linear_ramp tests
// ---------------------------------------------------------------------------

#[test]
fn test_pad_linear_ramp() {
    let cases: Vec<serde_json::Value> =
        serde_json::from_str(include_str!("references/beat/pad_linear_ramp.json")).unwrap();

    for case in &cases {
        let name = case["name"].as_str().unwrap();
        let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
        let pad_left = case["pad_left"].as_u64().unwrap() as usize;
        let pad_right = case["pad_right"].as_u64().unwrap() as usize;
        let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();

        let result = rosa::dsp::util::pad_linear_ramp(&input, pad_left, pad_right);
        assert_approx_eq(&result, &expected, 1e-12, name);
    }
}

// ---------------------------------------------------------------------------
// tempo_frequencies tests
// ---------------------------------------------------------------------------

#[test]
fn test_tempo_frequencies() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/tempo_frequencies.json")).unwrap();

    for case in data["tempo_frequencies"].as_array().unwrap() {
        let name = case["name"].as_str().unwrap();
        let n_bins = case["n_bins"].as_u64().unwrap() as usize;
        let hop_length = case["hop_length"].as_u64().unwrap() as usize;
        let sr = case["sr"].as_f64().unwrap();
        let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();

        let result = tempo_frequencies(n_bins, hop_length, sr);
        assert_approx_eq(
            &result,
            &expected,
            1e-10,
            &format!("tempo_frequencies_{name}"),
        );
    }
}

#[test]
fn test_fourier_tempo_frequencies() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/tempo_frequencies.json")).unwrap();

    for case in data["fourier_tempo_frequencies"].as_array().unwrap() {
        let name = case["name"].as_str().unwrap();
        let sr = case["sr"].as_f64().unwrap();
        let win_length = case["win_length"].as_u64().unwrap() as usize;
        let hop_length = case["hop_length"].as_u64().unwrap() as usize;
        let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();

        let result = fourier_tempo_frequencies(sr, win_length, hop_length);
        assert_approx_eq(
            &result,
            &expected,
            1e-10,
            &format!("fourier_tempo_frequencies_{name}"),
        );
    }
}

// ---------------------------------------------------------------------------
// ISTFT roundtrip test
// ---------------------------------------------------------------------------

#[test]
fn test_istft_roundtrip() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/istft.json")).unwrap();

    let original: Vec<f64> = serde_json::from_value(data["original"].clone()).unwrap();
    let n_fft = data["n_fft"].as_u64().unwrap() as usize;
    let hop_length = data["hop_length"].as_u64().unwrap() as usize;

    // Forward STFT
    let stft_params = rosa::stft::StftParams {
        n_fft,
        hop_length,
        win_length: n_fft,
        center: true,
        window: None,
    };
    let (stft_real, stft_imag) = rosa::stft::stft_as_real_imag(&original, &stft_params);

    // Inverse STFT
    let istft_params = rosa::stft::IstftParams {
        n_fft,
        hop_length,
        win_length: n_fft,
        center: true,
        length: None,
    };
    let reconstructed = rosa::stft::istft(&stft_real, &stft_imag, &istft_params);

    // Check roundtrip: stft -> istft should recover the original signal
    // in the core region (away from edges where zero-padding affects reconstruction)
    let margin = n_fft;
    let end = reconstructed
        .len()
        .min(original.len())
        .saturating_sub(margin);
    if end > margin {
        assert_approx_eq(
            &reconstructed[margin..end],
            &original[margin..end],
            1e-6,
            "istft_roundtrip_core",
        );
    }
}

// ---------------------------------------------------------------------------
// tempogram test
// ---------------------------------------------------------------------------

#[test]
fn test_tempogram() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/tempogram.json")).unwrap();

    let sr = data["sr"].as_f64().unwrap();
    let hop_length = data["hop_length"].as_u64().unwrap() as usize;
    let win_length = data["win_length"].as_u64().unwrap() as usize;
    let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
    let expected_shape: Vec<usize> = serde_json::from_value(data["shape"].clone()).unwrap();

    let params = TempogramParams {
        sr,
        hop_length,
        win_length,
        center: true,
        norm: Some(rosa::dsp::util::Norm::Linf),
    };
    let tg = tempogram(&onset_env, &params);

    // Check shape
    assert_eq!(tg.rows(), expected_shape[0], "tempogram rows mismatch");
    assert_eq!(tg.cols(), expected_shape[1], "tempogram cols mismatch");

    // Check sampled columns
    for col_data in data["sampled_columns"].as_array().unwrap() {
        let col_idx = col_data["index"].as_u64().unwrap() as usize;
        let expected_col: Vec<f64> = serde_json::from_value(col_data["data"].clone()).unwrap();

        if col_idx >= tg.cols() {
            continue;
        }

        let actual_col: Vec<f64> = (0..tg.rows()).map(|r| tg.get(r, col_idx)).collect();
        assert_approx_eq(
            &actual_col,
            &expected_col,
            1e-4,
            &format!("tempogram_col_{col_idx}"),
        );
    }
}

// ---------------------------------------------------------------------------
// tempo test
// ---------------------------------------------------------------------------

#[test]
fn test_tempo() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/tempo.json")).unwrap();

    let sr = data["sr"].as_f64().unwrap();
    let hop_length = data["hop_length"].as_u64().unwrap() as usize;
    let start_bpm = data["start_bpm"].as_f64().unwrap();
    let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
    let expected_tempo = data["tempo"].as_f64().unwrap();

    let params = TempoParams {
        sr,
        hop_length,
        start_bpm,
        std_bpm: 1.0,
        ac_size: 8.0,
        max_tempo: Some(320.0),
    };
    let result = tempo(&onset_env, &params);

    // Tempo should be reasonably close (within 5% or 5 BPM)
    let diff = (result - expected_tempo).abs();
    let rel_diff = diff / expected_tempo;
    assert!(
        rel_diff < 0.05 || diff < 5.0,
        "tempo: got {result}, expected {expected_tempo}, diff {diff} ({:.1}%)",
        rel_diff * 100.0
    );
}

// ---------------------------------------------------------------------------
// beat_track test
// ---------------------------------------------------------------------------

#[test]
fn test_beat_track() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/beat_track.json")).unwrap();

    let sr = data["sr"].as_f64().unwrap();
    let hop_length = data["hop_length"].as_u64().unwrap() as usize;
    let start_bpm = data["start_bpm"].as_f64().unwrap();
    let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
    let expected_tempo = data["tempo"].as_f64().unwrap();
    let expected_beats: Vec<usize> = serde_json::from_value(data["beats"].clone()).unwrap();

    let params = BeatTrackParams {
        sr,
        hop_length,
        start_bpm,
        tightness: 100.0,
        trim: true,
        bpm: None,
    };
    let (result_tempo, result_beats) = beat_track(&onset_env, &params);

    // Check tempo (within 10%)
    let tempo_diff = (result_tempo - expected_tempo).abs() / expected_tempo;
    assert!(
        tempo_diff < 0.1,
        "beat_track tempo: got {result_tempo}, expected {expected_tempo}"
    );

    // Check beats: allow some tolerance in positions
    // The number of beats should be similar (within 20%)
    if !expected_beats.is_empty() {
        let count_diff = (result_beats.len() as f64 - expected_beats.len() as f64).abs()
            / expected_beats.len() as f64;
        assert!(
            count_diff < 0.3,
            "beat_track beat count: got {}, expected {}",
            result_beats.len(),
            expected_beats.len()
        );
    }
}

// ---------------------------------------------------------------------------
// clicks beat_track test (deterministic)
// ---------------------------------------------------------------------------

#[test]
fn test_clicks_beat_track() {
    let data: serde_json::Value =
        serde_json::from_str(include_str!("references/beat/clicks_beat_track.json")).unwrap();

    let sr = data["sr"].as_f64().unwrap();
    let hop_length = data["hop_length"].as_u64().unwrap() as usize;
    let onset_env: Vec<f64> = serde_json::from_value(data["onset_envelope"].clone()).unwrap();
    let known_bpm = data["known_bpm"].as_f64().unwrap();

    let params = BeatTrackParams {
        sr,
        hop_length,
        start_bpm: 120.0,
        tightness: 100.0,
        trim: true,
        bpm: None,
    };
    let (result_tempo, result_beats) = beat_track(&onset_env, &params);

    // With synthetic clicks at 120 BPM, tempo should be close
    let tempo_diff = (result_tempo - known_bpm).abs();
    assert!(
        tempo_diff < 10.0,
        "clicks tempo: got {result_tempo}, expected ~{known_bpm}"
    );

    // Should detect some beats
    assert!(
        !result_beats.is_empty(),
        "clicks beat_track: should detect some beats"
    );
}