rosa 0.1.0

Audio feature extraction library matching librosa output
Documentation
use rosa::matrix::Matrix;
use rosa::sequence::*;

fn load_refs() -> serde_json::Value {
    let data = include_str!("references/sequence/sequence.json");
    serde_json::from_str(data).expect("Failed to parse sequence.json")
}

fn flatten_matrix(nested: &serde_json::Value) -> Vec<f64> {
    nested
        .as_array()
        .unwrap()
        .iter()
        .flat_map(|row| row.as_array().unwrap().iter().map(|v| v.as_f64().unwrap()))
        .collect()
}

fn assert_matrix_eq(actual: &Matrix, expected_nested: &serde_json::Value, tol: f64, name: &str) {
    let expected = flatten_matrix(expected_nested);
    let actual_data = actual.as_slice();
    assert_eq!(
        actual_data.len(),
        expected.len(),
        "{name}: size mismatch: {} vs {}",
        actual_data.len(),
        expected.len()
    );
    for (i, (a, e)) in actual_data.iter().zip(expected.iter()).enumerate() {
        let diff = (a - e).abs();
        assert!(
            diff < tol,
            "{name}[{i}]: got {a}, expected {e}, diff {diff}"
        );
    }
}

// --- Transition matrix tests ---

#[test]
fn test_transition_uniform_3_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_uniform_3"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let t = transition_uniform(n);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_uniform_3");
}

#[test]
fn test_transition_uniform_5_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_uniform_5"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let t = transition_uniform(n);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_uniform_5");
}

#[test]
fn test_transition_loop_scalar_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_loop_3_0.5"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let prob: Vec<f64> = serde_json::from_value(case["prob"].clone()).unwrap();
    let t = transition_loop(n, &prob);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_loop_3_0.5");
}

#[test]
fn test_transition_loop_varied_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_loop_3_varied"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let prob: Vec<f64> = serde_json::from_value(case["prob"].clone()).unwrap();
    let t = transition_loop(n, &prob);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_loop_3_varied");
}

#[test]
fn test_transition_loop_4_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_loop_4_0.9"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let prob: Vec<f64> = serde_json::from_value(case["prob"].clone()).unwrap();
    let t = transition_loop(n, &prob);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_loop_4_0.9");
}

#[test]
fn test_transition_local_5_3_triangle_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_local_5_3_triangle"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let w = case["width"].as_u64().unwrap() as usize;
    let t = transition_local(n, w, "triangle", false);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_local_5_3_triangle");
}

#[test]
fn test_transition_local_5_3_triangle_wrap_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_local_5_3_triangle_wrap"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let w = case["width"].as_u64().unwrap() as usize;
    let t = transition_local(n, w, "triangle", true);
    assert_matrix_eq(
        &t,
        &case["matrix"],
        1e-12,
        "transition_local_5_3_triangle_wrap",
    );
}

#[test]
fn test_transition_local_5_3_ones_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_local_5_3_ones"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let w = case["width"].as_u64().unwrap() as usize;
    let t = transition_local(n, w, "ones", false);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_local_5_3_ones");
}

#[test]
fn test_transition_local_8_5_triangle_matches_librosa() {
    let refs = load_refs();
    let case = &refs["transition_local_8_5_triangle"];
    let n = case["n_states"].as_u64().unwrap() as usize;
    let w = case["width"].as_u64().unwrap() as usize;
    let t = transition_local(n, w, "triangle", false);
    assert_matrix_eq(&t, &case["matrix"], 1e-12, "transition_local_8_5_triangle");
}

// --- Viterbi tests ---

#[test]
fn test_viterbi_2state_matches_librosa() {
    let refs = load_refs();
    let case = &refs["viterbi_2state"];

    let prob_flat = flatten_matrix(&case["prob"]);
    let n_states = case["prob"].as_array().unwrap().len();
    let n_steps = case["prob"][0].as_array().unwrap().len();
    let prob = Matrix::from_vec(prob_flat, n_states, n_steps);

    let trans_flat = flatten_matrix(&case["transition"]);
    let trans = Matrix::from_vec(trans_flat, n_states, n_states);

    let p_init: Vec<f64> = serde_json::from_value(case["p_init"].clone()).unwrap();

    let path = viterbi(&prob, &trans, Some(&p_init));
    let expected: Vec<usize> = serde_json::from_value(case["path"].clone()).unwrap();
    assert_eq!(path, expected, "viterbi_2state path mismatch");
}

#[test]
fn test_viterbi_3state_matches_librosa() {
    let refs = load_refs();
    let case = &refs["viterbi_3state"];

    let prob_flat = flatten_matrix(&case["prob"]);
    let n_states = case["prob"].as_array().unwrap().len();
    let n_steps = case["prob"][0].as_array().unwrap().len();
    let prob = Matrix::from_vec(prob_flat, n_states, n_steps);

    let trans_flat = flatten_matrix(&case["transition"]);
    let trans = Matrix::from_vec(trans_flat, n_states, n_states);

    let path = viterbi(&prob, &trans, None);
    let expected: Vec<usize> = serde_json::from_value(case["path"].clone()).unwrap();
    assert_eq!(path, expected, "viterbi_3state path mismatch");
}

#[test]
fn test_viterbi_discriminative_matches_librosa() {
    let refs = load_refs();
    let case = &refs["viterbi_discriminative_2state"];

    let prob_flat = flatten_matrix(&case["prob"]);
    let n_states = case["prob"].as_array().unwrap().len();
    let n_steps = case["prob"][0].as_array().unwrap().len();
    let prob = Matrix::from_vec(prob_flat, n_states, n_steps);

    let trans_flat = flatten_matrix(&case["transition"]);
    let trans = Matrix::from_vec(trans_flat, n_states, n_states);

    let path = viterbi_discriminative(&prob, &trans, None, None);
    let expected: Vec<usize> = serde_json::from_value(case["path"].clone()).unwrap();
    assert_eq!(path, expected, "viterbi_discriminative path mismatch");
}

#[test]
fn test_viterbi_binary_matches_librosa() {
    let refs = load_refs();
    let case = &refs["viterbi_binary"];

    let prob_flat = flatten_matrix(&case["prob"]);
    let n_states = case["prob"].as_array().unwrap().len();
    let n_steps = case["prob"][0].as_array().unwrap().len();
    let prob = Matrix::from_vec(prob_flat, n_states, n_steps);

    let trans_flat = flatten_matrix(&case["transition"]);
    let trans = Matrix::from_vec(trans_flat, 2, 2);

    let result = viterbi_binary(&prob, &trans, Some(0.5), Some(0.5));
    let expected_flat = flatten_matrix(&case["states"]);
    let expected_n_states = case["states"].as_array().unwrap().len();
    let expected_n_steps = case["states"][0].as_array().unwrap().len();

    assert_eq!(result.shape(), (expected_n_states, expected_n_steps));
    for (i, (a, e)) in result
        .as_slice()
        .iter()
        .zip(expected_flat.iter())
        .enumerate()
    {
        assert_eq!(
            *a as usize, *e as usize,
            "viterbi_binary[{i}]: got {a}, expected {e}"
        );
    }
}